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

const std = @import("std");

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

pub var trap_frame: trap.Frame = undefined;

pub const SupervisorTrapVector = packed struct(usize) {
    pub const Mode = enum(u2) {
        direct = 0,
        vectored = 1,
    };

    mode: u2,
    base_addr: u62,

    pub fn fromBaseAddr(mode: Mode, base_addr: usize) SupervisorTrapVector {
        return .{
            .mode = @intFromEnum(mode),
            .base_addr = base_addr >> 2,
        };
    }
};

pub const Enable = packed struct(usize) {
    u_software: u1,
    s_software: u1,
    reserved0: u2,
    u_timer: u1,
    s_timer: u1,
    reserved1: u2,
    u_external: u1,
    s_external: u1,
    reserved2: u54,

    pub const none = std.mem.zeroInit(Enable, .{});
    pub const all = std.mem.zeroInit(Enable, .{
        .u_software = 1,
        .s_software = 1,
        .u_timer = 1,
        .s_timer = 1,
        .u_external = 1,
        .s_external = 1,
    });
};

pub const Cause = packed struct(usize) {
    num: u63,
    @"async": u1,

    pub fn isAsync(self: Cause) bool {
        return self.@"async" == 1;
    }
};

pub const AsyncCause = enum(u63) {
    user_software = 0,
    supervisor_software = 1,
    user_timer = 4,
    supervisor_timer = 5,
    user_external = 8,
    supervisor_external = 9,
    _,
};

pub const SyncCause = enum(u63) {
    instruction_address_misaligned = 0,
    instruction_access_fault = 1,
    illegal_instruction = 2,
    breakpoint = 3,
    load_access_fault = 5,
    amo_address_misaligned = 6,
    store_or_amo_access_fault = 7,
    ecall = 8,
    instruction_page_fault = 12,
    load_page_fault = 13,
    store_or_amo_page_fault = 15,
    _,
};

export fn handleTrap(epc: usize, tval: usize, cause_bits: usize, frame: *trap.Frame) usize {
    const console = Console.autoChoose().?;
    const w = console.writer();

    const status = instructions.sstatus.read();

    const cause: Cause = @bitCast(cause_bits);

    if (cause.isAsync()) {
        switch (@as(AsyncCause, @enumFromInt(cause.num))) {
            .supervisor_software => w.print("Hart {d}: Software interrupt\r\n", .{frame.hart_id}) catch unreachable,
            .supervisor_timer => {
                defer time.interruptInMillis(process.schedule_interval_millis) catch |err| {
                    std.debug.panic("Hart {d}: Unable to reset interrupt timer: {any}", .{ frame.hart_id, err });
                };

                if (status.previous_privilege == .user) {
                    // Trapped from U-mode, update pc for next time slice.
                    //
                    // We can simply use the last node of the process list here
                    // because the scheduler moves a process to the end of the queue
                    // before returning into it.

                    process.list.last.?.data.pc = epc;
                    process.list.last.?.data.state = .waiting;

                    process.schedule() catch |err| {
                        std.debug.panic("Hart {d}: Unable to schedule next process: {any}", .{ frame.hart_id, err });
                    };
                }

                // Don't interrupt kernel code, it may never run otherwise.
            },
            .supervisor_external => {
                const context: u14 = @intCast(2 * frame.hart_id + 1);

                const external_cause = plic.default.claim(context) catch |err| {
                    std.debug.panic("Hart {d}: Unable to claim external interrupt: {any}", .{ frame.hart_id, err });
                };
                if (external_cause) |source| {
                    w.print("Hart {d}: External interrupt: 0x{x}\r\n", .{ frame.hart_id, source }) catch unreachable;
                    handleExternal(source);
                    plic.default.complete(context, source) catch |err| {
                        std.debug.panic("Hart {d}: Unable to complete external interrupt: {any}", .{ frame.hart_id, err });
                    };
                }
            },
            else => {
                std.debug.panic("Hart {d}: Unhandled asynchronous interrupt: {d}", .{ frame.hart_id, cause.num });
            },
        }
    } else {
        const pid = if (status.previous_privilege == .user) blk: {
            const proc = &process.list.last.?.data;
            proc.pc = epc;
            proc.state = .waiting;
            break :blk proc.id;
        } else 0;

        switch (@as(SyncCause, @enumFromInt(cause.num))) {
            .illegal_instruction => {
                std.debug.panic("Hart {d}, PID = {d}: Illegal instruction, EPC = 0x{x:0>16}", .{ frame.hart_id, pid, epc });
            },
            .instruction_access_fault => {
                std.debug.panic("Hart {d}, PID = {d}: Instruction access fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}", .{ frame.hart_id, pid, epc, tval });
            },
            .load_access_fault => {
                std.debug.panic("Hart {d}, PID = {d}: Load access fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}", .{ frame.hart_id, pid, epc, tval });
            },
            .store_or_amo_access_fault => {
                std.debug.panic("Hart {d}, PID = {d}: Store/AMO access fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}", .{ frame.hart_id, pid, epc, tval });
            },
            .ecall => {
                const proc = &process.list.last.?.data;
                syscall.handler(proc, frame) catch |err| switch (err) {
                    syscall.HandleError.UnknownSyscall => {
                        const a7 = frame.general_purpose_registers[17];
                        w.print("Hart {d}, PID = {d}: Unknown syscall, a7 = 0x{x:0>16}\r\n", .{ frame.hart_id, pid, a7 }) catch unreachable;
                    },
                };

                return epc + 4;
            },
            .instruction_page_fault => {
                std.debug.panic("Hart {d}, PID = {d}: Instruction page fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}", .{ frame.hart_id, pid, epc, tval });
            },
            .load_page_fault => {
                std.debug.panic("Hart {d}, PID = {d}: Load page fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}", .{ frame.hart_id, pid, epc, tval });
            },
            .store_or_amo_page_fault => {
                std.debug.panic("Hart {d}, PID = {d}: Store/AMO page fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}", .{ frame.hart_id, pid, epc, tval });
            },
            else => {
                std.debug.panic("Hart {d}, PID = {d}: Unhandled synchronous interrupt: {d}, EPC = 0x{x:0>16}", .{ frame.hart_id, pid, cause.num, epc });
            },
        }
    }

    return epc;
}

fn handleExternal(interrupt: ?u10) void {
    _ = &interrupt;
}

export fn supervisorTrapVector() align(4) callconv(.Naked) noreturn {
    asm volatile (
        \\ csrrw t6, sscratch, t6
        \\
        \\ sd x1, 8(t6)
        \\ sd x2, 16(t6)
        \\ sd x3, 24(t6)
        \\ sd x4, 32(t6)
        \\ sd x5, 40(t6)
        \\ sd x6, 48(t6)
        \\ sd x7, 56(t6)
        \\ sd x8, 64(t6)
        \\ sd x9, 72(t6)
        \\ sd x10, 80(t6)
        \\ sd x11, 88(t6)
        \\ sd x12, 96(t6)
        \\ sd x13, 104(t6)
        \\ sd x14, 112(t6)
        \\ sd x15, 120(t6)
        \\ sd x16, 128(t6)
        \\ sd x17, 136(t6)
        \\ sd x18, 144(t6)
        \\ sd x19, 152(t6)
        \\ sd x20, 160(t6)
        \\ sd x21, 168(t6)
        \\ sd x22, 176(t6)
        \\ sd x23, 184(t6)
        \\ sd x24, 192(t6)
        \\ sd x25, 200(t6)
        \\ sd x26, 208(t6)
        \\ sd x27, 216(t6)
        \\ sd x28, 224(t6)
        \\ sd x29, 232(t6)
        \\ sd x30, 240(t6)
        \\
        \\ mv t5, t6
        \\ csrr t6, sscratch
        \\
        \\ sd x31, 248(t5)
        \\
        \\ csrw sscratch, t5
        \\
        \\ .option push
        \\ .option norelax
        \\ la gp, _global_pointer
        \\ .option pop
        \\
        \\ csrr a0, sepc
        \\ csrr a1, stval
        \\ csrr a2, scause
        \\ mv a3, t5
        \\
        \\ la sp, _stvec_stack_end
        \\ call handleTrap
        \\
        \\ csrw sepc, a0
        \\
        \\ 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
    );
}

pub fn init(hart_id: usize) void {
    trap_frame = .{
        .general_purpose_registers = [_]usize{0} ** 32,
        .floating_point_registers = [_]usize{0} ** 32,
        .satp = 0,
        .stack_pointer = @ptrFromInt(instructions.stackPointer()),
        .hart_id = hart_id,
    };

    asm volatile (
        \\ csrw sscratch, %[trapframe]
        \\
        \\ la t0, supervisorTrapVector
        \\ csrw stvec, t0
        \\
        \\ csrr t0, sstatus
        \\ ori t0, t0, 2
        \\ csrw sstatus, t0
        :
        : [trapframe] "r" (&trap_frame),
    );
}