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

pub const ProcessError = error{
    TooManyResourceDescriptors,
};

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

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

pub const Info = struct {
    id: u16,
    trap_frame: trap.Frame,
    stack: *[num_stack_pages * paging.page_size]u8,
    pc: usize,
    page_table: *paging.Table,
    state: State,
    rds: std.AutoArrayHashMap(usize, vfs.ResourceDescriptor),

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

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

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

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

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

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;

    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 != phdr.p_memsz) return ExeError.SizeMismatch;
        if (phdr.p_filesz > elf_buf.len or phdr.p_memsz > pages.len) return ExeError.MemOverrun;

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

        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;
            }

            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{
        .id = next_pid,
        .trap_frame = std.mem.zeroInit(trap.Frame, .{}),
        .stack = @ptrCast(stack),
        .pc = hdr.entry,
        .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 Error.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;
}