// SPDX-FileCopyrightText: 2024 Himbeer <himbeer@disroot.org>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

const std = @import("std");

const instructions = @import("instructions.zig");
const paging = @import("paging.zig");
const time = @import("sbi/time.zig");
const trap = @import("trap.zig");

pub const schedule_interval_millis = 1;

pub var list = std.mem.zeroInit(std.DoublyLinkedList(Info), .{});

const num_stack_pages = 2;

var next_pid: u16 = 1;

pub const Error = error{
    EmptySchedule,
};

pub const State = enum(u8) {
    waiting,
    active,
    sleeping,
    terminated,
};

pub const Info = extern struct {
    id: u16,
    trap_frame: trap.Frame,
    stack: [*]u8,
    pc: usize,
    page_table: *paging.Table,
    state: State,

    pub fn destroy(self: *Info) !void {
        try paging.free(self.stack);
        try self.page_table.unmap();
        try paging.free(self.page_table);
    }

    pub fn satp(self: *const Info) paging.Satp {
        return self.page_table.satp(self.id);
    }
};

fn new(entry: usize) !Info {
    const stack = try paging.alloc(num_stack_pages);
    errdefer paging.free(stack) catch {};

    const procmem: *paging.Table = @alignCast(@ptrCast(try paging.zeroedAlloc(1)));
    errdefer paging.free(@ptrCast(procmem)) catch {};

    var proc = Info{
        .id = next_pid,
        .trap_frame = std.mem.zeroInit(trap.Frame, .{}),
        .stack = @alignCast(@ptrCast(stack)),
        .pc = entry,
        .page_table = procmem,
        .state = .waiting,
    };

    const stack_top = @intFromPtr(proc.stack) + num_stack_pages * paging.page_size;
    proc.trap_frame.general_purpose_registers[2] = stack_top;

    try procmem.mapKernel();

    try procmem.map(entry, entry, paging.EntryFlags.userReadExec, 0);
    // Not using identityMapRange because this is going to be expanded for non-relocatable binaries.
    for (0..num_stack_pages) |page| {
        const vaddr = @intFromPtr(proc.stack) + page * paging.page_size;
        const paddr = @intFromPtr(proc.stack) + page * paging.page_size;

        try procmem.map(vaddr, paddr, paging.EntryFlags.userReadWrite, 0);
    }

    next_pid += 1;
    return proc;
}

pub fn next() ?*Info {
    if (list.popFirst()) |info| {
        list.append(info);
        return &info.data;
    }

    return null;
}

pub fn switchTo(proc: *Info) noreturn {
    proc.state = .active;

    instructions.sscratch.write(@intFromPtr(&proc.trap_frame));

    var sstatus = instructions.sstatus.read();
    sstatus.previous_privilege = .user;
    instructions.sstatus.write(sstatus);

    instructions.sepc.write(proc.pc);
    instructions.satp.write(proc.satp());

    // Probably not always needed. Let's not take the risk for now.
    asm volatile (
        \\ sfence.vma
    );

    asm volatile (
        \\ csrr t6, sscratch
        \\
        \\ ld x1, 8(t6)
        \\ ld x2, 16(t6)
        \\ ld x3, 24(t6)
        \\ ld x4, 32(t6)
        \\ ld x5, 40(t6)
        \\ ld x6, 48(t6)
        \\ ld x7, 56(t6)
        \\ ld x8, 64(t6)
        \\ ld x9, 72(t6)
        \\ ld x10, 80(t6)
        \\ ld x11, 88(t6)
        \\ ld x12, 96(t6)
        \\ ld x13, 104(t6)
        \\ ld x14, 112(t6)
        \\ ld x15, 120(t6)
        \\ ld x16, 128(t6)
        \\ ld x17, 136(t6)
        \\ ld x18, 144(t6)
        \\ ld x19, 152(t6)
        \\ ld x20, 160(t6)
        \\ ld x21, 168(t6)
        \\ ld x22, 176(t6)
        \\ ld x23, 184(t6)
        \\ ld x24, 192(t6)
        \\ ld x25, 200(t6)
        \\ ld x26, 208(t6)
        \\ ld x27, 216(t6)
        \\ ld x28, 224(t6)
        \\ ld x29, 232(t6)
        \\ ld x30, 240(t6)
        \\ ld x31, 248(t6)
        \\
        \\ sret
    );

    unreachable;
}

pub fn demo(allocator: std.mem.Allocator) !noreturn {
    const entry: [*]u8 = @alignCast(@ptrCast(try paging.zeroedAlloc(1)));
    defer paging.free(@ptrCast(entry)) catch {};

    entry[0] = 0x73;
    entry[1] = 0x00;
    entry[2] = 0x00;
    entry[3] = 0x00;

    const proc = try new(@intFromPtr(entry));
    const proc_node = try allocator.create(std.DoublyLinkedList(Info).Node);
    proc_node.data = proc;
    list.prepend(proc_node);

    try time.interruptInMillis(schedule_interval_millis);
    try switchTo(&proc_node.data);
}