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

const builtin = @import("builtin");
const std = @import("std");

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

const elf = std.elf;

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 ScheduleError = error{
    EmptySchedule,
    NoInit,
};

pub const Error = error{
    TooManyThreads,
    TooManyResourceDescriptors,
    BadRdHandle,
};

pub const ExeError = error{
    BadEndian,
    BadArch,
    BadBitLen,
    NotStaticExe,
    SizeMismatch,
    MemOverrun,
    BranchPerms,
    WritableCode,
};

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

pub const Info = struct {
    allocator: std.mem.Allocator,
    id: u16,
    thread_id: usize,
    trap_frame: trap.Frame,
    pages: []align(paging.page_size) u8,
    stack: []align(paging.page_size) u8,
    pc: usize,
    cleanup_hook: ?CleanupHook,
    term_hook: ?TermHook,
    page_table: *paging.Table,
    state: State,
    rds: std.AutoArrayHashMap(usize, vfs.ResourceDescriptor),

    pub const CleanupHook = struct {
        pub const CleanupFn = *const fn (proc: *const Info, buffer: []u8, copy: []const u8) void;

        cleanupFn: CleanupFn,
        buffer: []u8,
        copy: []const u8,
    };

    pub const TermHook = struct {
        pub const HookFn = *const fn (context: *anyopaque, proc: *const Info) void;

        hookFn: HookFn,
        context: *anyopaque,
    };

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

    pub fn createRdHandle(self: *Info, rd: vfs.ResourceDescriptor) !usize {
        if (self.rds.count() == std.math.maxInt(usize)) {
            return Error.TooManyResourceDescriptors;
        }

        var handle = self.rds.count() + 1;
        while (self.rds.contains(handle)) : (handle +%= 1) {}

        try self.rds.put(handle, rd);
        return handle;
    }

    pub fn destroyRdHandle(self: *Info, handle: usize) void {
        if (self.rds.fetchSwapRemove(handle)) |kv| {
            kv.value.deinit();
        }
    }

    pub fn createThread(self: *const Info, allocator: ?std.mem.Allocator) !*Info {
        const alloc = allocator orelse self.allocator;

        var trap_frame = std.mem.zeroInit(trap.Frame, .{});

        const stack = try paging.zeroedAlloc(num_stack_pages);
        errdefer paging.free(stack);

        const stack_top = @intFromPtr(stack.ptr) + num_stack_pages * paging.page_size;
        try self.page_table.identityMapRange(@intFromPtr(stack.ptr), stack_top, paging.EntryFlags.userReadWrite);

        trap_frame.general_purpose_registers[2] = stack_top;

        const thread_id = std.math.add(usize, self.thread_id, 1) catch {
            return Error.TooManyThreads;
        };

        const proc = .{
            .allocator = alloc,
            .id = self.id,
            .thread_id = thread_id,
            .trap_frame = trap_frame,
            .pages = self.pages,
            .stack = stack,
            .pc = self.pc,
            .cleanup_hook = null,
            .term_hook = null,
            .page_table = self.page_table,
            .state = .suspended,
            .rds = self.rds,
        };

        const proc_node = try alloc.create(std.DoublyLinkedList(Info).Node);
        proc_node.data = proc;
        list.prepend(proc_node);

        return &proc_node.data;
    }

    pub fn call(self: *Info, function: usize, args: anytype, cleanup_hook: ?CleanupHook, term_hook: ?TermHook) noreturn {
        self.pc = function;
        self.cleanup_hook = cleanup_hook;
        self.term_hook = term_hook;
        self.trap_frame.general_purpose_registers[1] = @intFromPtr(&rethooks.terminate);
        inline for (args, 0..) |arg, i| {
            if (i >= 6) break;

            const value = usizeFromArg(arg);
            self.trap_frame.general_purpose_registers[10 + i] = value;
        }

        self.state = .waiting;
        schedule() catch |err| {
            std.debug.panic("Unable to schedule thread: {any}", .{err});
        };
    }

    pub fn terminate(
        self: *Info,
    ) void {
        var node = list.first;
        while (node) |proc_node| : (node = proc_node.next) {
            if (self.shouldRemove(&proc_node.data)) {
                if (proc_node.data.thread_id != self.thread_id) {
                    proc_node.data.terminate();
                }

                list.remove(proc_node);
                self.allocator.destroy(proc_node);
            }
        }

        paging.free(self.stack);

        if (self.thread_id == 0) {
            self.page_table.unmap();
            paging.free(self.page_table);

            paging.free(self.pages);

            for (self.rds.values()) |rd| {
                rd.deinit();
            }
            self.rds.clearAndFree();
        }

        if (self.cleanup_hook) |cleanup_hook| {
            cleanup_hook.cleanupFn(self, cleanup_hook.buffer, cleanup_hook.copy);
        }
        if (self.term_hook) |term_hook| {
            sysexchange.frameReturn(void, &self.trap_frame, vfs.Error.Detached);
            term_hook.hookFn(term_hook.context, self);
        }
    }

    pub fn allowResume(self: *Info) void {
        self.pc += 4; // Skip ecall instruction
        self.state = .waiting;
    }

    fn shouldRemove(self: *const Info, candidate: *const Info) bool {
        return candidate.id == self.id and self.shouldRemoveThread(candidate);
    }

    fn shouldRemoveThread(self: *const Info, candidate: *const Info) bool {
        return candidate.thread_id == self.thread_id or self.thread_id == 0;
    }

    pub fn copyBytes(self: *const Info, bytes: []const u8) ![]align(paging.page_size) const u8 {
        const aligned_len = std.mem.alignForward(usize, bytes.len, paging.page_size);
        const num_pages = @divExact(aligned_len, paging.page_size);

        const copy = try paging.zeroedAlloc(num_pages);
        errdefer paging.free(copy);

        var addr = @intFromPtr(copy.ptr);
        const limit = addr + copy.len;
        while (addr < limit) : (addr += paging.page_size) {
            try self.page_table.map(addr, addr, paging.EntryFlags.userReadOnly, 0);
            errdefer self.page_table.unmapEntry(addr);
        }

        paging.setUserMemoryAccess(true);
        defer paging.setUserMemoryAccess(false);

        @memcpy(copy[0..bytes.len], bytes);
        return copy[0..bytes.len];
    }

    pub fn copyBuffer(self: *const Info, buffer: []u8) ![]align(paging.page_size) u8 {
        const aligned_len = std.mem.alignForward(usize, buffer.len, paging.page_size);
        const num_pages = @divExact(aligned_len, paging.page_size);

        const copy = try paging.zeroedAlloc(num_pages);
        errdefer paging.free(copy);

        var addr = @intFromPtr(copy.ptr);
        const limit = addr + copy.len;
        while (addr < limit) : (addr += paging.page_size) {
            try self.page_table.map(addr, addr, paging.EntryFlags.userReadWrite, 0);
            errdefer self.page_table.unmapEntry(addr);
        }

        paging.setUserMemoryAccess(true);
        defer paging.setUserMemoryAccess(false);

        @memcpy(copy[0..buffer.len], buffer);
        return copy[0..buffer.len];
    }
};

pub fn next() ?*Info {
    if (list.popFirst()) |info| {
        list.append(info);
        // fixme: Suspending or sleeping init process causes infinite recursion.
        if (info.data.state != .waiting) return next();
        return &info.data;
    }

    return null;
}

pub fn schedule() !noreturn {
    if (next()) |proc| {
        try time.interruptInMillis(schedule_interval_millis);
        switchTo(proc);
    }

    return ScheduleError.EmptySchedule;
}

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

    var sstatus = instructions.sstatus.read();
    sstatus.previous_privilege = .user;
    sstatus.user_interrupt_enable = 0;
    sstatus.supervisor_interrupt_enable = 0;
    sstatus.user_prior_interrupt_enable = 1;
    sstatus.supervisor_prior_interrupt_enable = 1;
    instructions.sstatus.write(sstatus);

    instructions.sscratch.write(@intFromPtr(&proc.trap_frame));
    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 create(allocator: std.mem.Allocator, elf_buf: []align(@alignOf(elf.Elf64_Ehdr)) const u8) !*Info {
    const hdr_buf: *align(@alignOf(elf.Elf64_Ehdr)) const [@sizeOf(elf.Elf64_Ehdr)]u8 = elf_buf[0..@sizeOf(elf.Elf64_Ehdr)];
    const hdr = try elf.Header.parse(@ptrCast(hdr_buf));

    try validateElfHeader(hdr, hdr_buf);

    const len_aligned = std.mem.alignForwardLog2(elf_buf.len, paging.log2_page_size);
    const num_pages = len_aligned / paging.page_size + 1;
    const pages = try paging.zeroedAlloc(num_pages);
    errdefer paging.free(pages);

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

    try procmem.mapKernel();

    const parse_source = std.io.fixedBufferStream(elf_buf);

    var it = hdr.program_header_iterator(parse_source);
    while (try it.next()) |phdr| {
        if (phdr.p_type != elf.PT_LOAD) continue;
        if (phdr.p_filesz == 0 or phdr.p_memsz == 0) continue;
        if (phdr.p_filesz > elf_buf.len or phdr.p_memsz > pages.len) return ExeError.MemOverrun;

        const sz = @min(phdr.p_filesz, phdr.p_memsz);
        @memcpy(pages[phdr.p_offset .. phdr.p_offset + sz], elf_buf[phdr.p_offset .. phdr.p_offset + sz]);

        const memsz_aligned = std.mem.alignForwardLog2(phdr.p_memsz, paging.log2_page_size);
        const num_mappings = @divExact(memsz_aligned, paging.page_size);
        for (0..num_mappings) |page| {
            const vaddr = phdr.p_vaddr + page * paging.page_size;
            const paddr = @intFromPtr(pages.ptr) + phdr.p_offset + page * paging.page_size;
            const flags = paging.EntryFlags{
                .valid = 1,
                .read = @bitCast(phdr.p_flags & elf.PF_R != 0),
                .write = @bitCast(phdr.p_flags & elf.PF_W != 0),
                .exec = @bitCast(phdr.p_flags & elf.PF_X != 0),
                .user = 1,
                .global = 0,
                .accessed = 1,
                .dirty = @bitCast(phdr.p_flags & elf.PF_W != 0),
            };

            if (!@bitCast(flags.read) and !@bitCast(flags.write) and !@bitCast(flags.exec)) {
                return ExeError.BranchPerms;
            }
            if (@bitCast(flags.write) and @bitCast(flags.exec)) {
                return ExeError.WritableCode;
            }

            try procmem.map(vaddr, paddr, flags, 0);
        }
    }

    const stack = try paging.zeroedAlloc(num_stack_pages);
    errdefer paging.free(stack);

    const stack_top = @intFromPtr(stack.ptr) + num_stack_pages * paging.page_size;
    try procmem.identityMapRange(@intFromPtr(stack.ptr), stack_top, paging.EntryFlags.userReadWrite);

    var proc = Info{
        .allocator = allocator,
        .id = next_pid,
        .thread_id = 0,
        .trap_frame = std.mem.zeroInit(trap.Frame, .{}),
        .pages = pages,
        .stack = @ptrCast(stack),
        .pc = hdr.entry,
        .cleanup_hook = null,
        .term_hook = null,
        .page_table = procmem,
        .state = .waiting,
        .rds = std.AutoArrayHashMap(usize, vfs.ResourceDescriptor).init(allocator),
    };
    proc.trap_frame.general_purpose_registers[2] = stack_top;

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

    return &proc_node.data;
}

pub fn runInit(allocator: std.mem.Allocator, reader: anytype) !noreturn {
    var file_name_buffer: [4096]u8 = undefined;
    var link_name_buffer: [4096]u8 = undefined;

    var it = std.tar.iterator(reader, .{
        .file_name_buffer = file_name_buffer[0..],
        .link_name_buffer = link_name_buffer[0..],
    });
    const exe = while (try it.next()) |file| {
        if (std.mem.eql(u8, file.name, "./init")) {
            break file;
        }
    } else return ScheduleError.NoInit;

    const alignment = @alignOf(elf.Elf64_Ehdr);

    var exe_list = std.ArrayListAligned(u8, alignment).init(allocator);
    defer exe_list.deinit();
    try exe.reader().readAllArrayListAligned(alignment, &exe_list, exe.size);

    const proc = try create(allocator, exe_list.items);

    try time.interruptInMillis(schedule_interval_millis);
    switchTo(proc);
}

fn validateElfHeader(hdr: elf.Header, hdr_buf: *align(@alignOf(elf.Elf64_Ehdr)) const [@sizeOf(elf.Elf64_Ehdr)]u8) !void {
    const arch = builtin.cpu.arch;

    if (hdr.endian != arch.endian()) return ExeError.BadEndian;
    if (hdr.machine != arch.toElfMachine()) return ExeError.BadArch;
    if (!hdr.is_64) return ExeError.BadBitLen;

    const hdr64 = @as(*const elf.Elf64_Ehdr, @ptrCast(hdr_buf));

    if (hdr64.e_type != .EXEC) return ExeError.NotStaticExe;
}

fn usizeFromArg(arg: anytype) usize {
    return switch (@typeInfo(@TypeOf(arg))) {
        .Pointer => |ptr| switch (ptr.size) {
            .Slice => @intFromPtr(arg.ptr),
            else => @intFromPtr(arg),
        },
        else => arg,
    };
}

pub fn findThread(pid: u16, thread_id: usize) ?*Info {
    var node = list.first;
    while (node) |proc_node| : (node = proc_node.next) {
        if (proc_node.data.id == pid and proc_node.data.thread_id == thread_id) {
            return &proc_node.data;
        }
    }

    return null;
}

pub fn mainThread(pid: u16) ?*Info {
    return findThread(pid, 0);
}

pub fn latestThread(pid: u16) ?*Info {
    var latest: ?*Info = null;

    var node = list.first;
    while (node) |proc_node| : (node = proc_node.next) {
        if (proc_node.data.id == pid) {
            if (latest) |proc| {
                if (proc_node.data.thread_id > proc.thread_id) {
                    latest = &proc_node.data;
                }
            } else latest = &proc_node.data;
        }
    }

    return latest;
}