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

const std = @import("std");

const debug_console = @import("sbi/debug_console.zig");
const instructions = @import("instructions.zig");

pub var trap_frame: TrapFrame = 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 all = .{
        .u_software = 1,
        .s_software = 1,
        .reserved0 = 0,
        .u_timer = 1,
        .s_timer = 1,
        .reserved1 = 0,
        .u_external = 1,
        .s_external = 1,
        .reserved2 = 0,
    };
};

pub const TrapFrame = extern struct {
    general_purpose_registers: [32]usize, // Offset: 0
    floating_point_registers: [32]usize, // Offset: 256
    satp: usize, // Offset: 512
    stack_pointer: *u8, // Offset: 520
    hart_id: usize, // Offset: 528
};

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

    pub fn is_async(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 supervisor_trap(epc: usize, tval: usize, cause_bits: usize, hart_id: usize, status: usize, frame: *TrapFrame) usize {
    _ = &status;
    _ = &frame;

    const w = debug_console.writer() catch while (true) {};

    const cause: Cause = @bitCast(cause_bits);

    if (cause.is_async()) {
        switch (@as(AsyncCause, @enumFromInt(cause.num))) {
            .supervisor_software => w.print("Hart {d}: Software interrupt\r\n", .{hart_id}) catch while (true) {},
            .supervisor_timer => w.print("Hart {d}: Timer interrupt\r\n", .{hart_id}) catch while (true) {},
            .supervisor_external => w.print("Hart {d}: External interrupt\r\n", .{hart_id}) catch while (true) {},
            else => {
                w.print("Hart {d}: Unhandled asynchronous interrupt: {d}\r\n", .{ hart_id, cause.num }) catch while (true) {};
                while (true) {}
            },
        }
    } else {
        switch (@as(SyncCause, @enumFromInt(cause.num))) {
            .illegal_instruction => {
                w.print("Hart {d}: Illegal instruction\r\n", .{hart_id}) catch while (true) {};
                while (true) {}
            },
            .instruction_access_fault => {
                w.print("Hart {d}: Instruction access fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}\r\n", .{ hart_id, epc, tval }) catch while (true) {};
                return epc + 4;
            },
            .load_access_fault => {
                w.print("Hart {d}: Load access fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}\r\n", .{ hart_id, epc, tval }) catch while (true) {};
                return epc + 4;
            },
            .store_or_amo_access_fault => {
                w.print("Hart {d}: Store/AMO access fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}\r\n", .{ hart_id, epc, tval }) catch while (true) {};
                return epc + 4;
            },
            .ecall => {
                w.print("Hart {d}: Environment call from U-mode\r\n", .{hart_id}) catch while (true) {};
                return epc + 4;
            },
            .instruction_page_fault => {
                w.print("Hart {d}: Instruction page fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}\r\n", .{ hart_id, epc, tval }) catch while (true) {};
                return epc + 4;
            },
            .load_page_fault => {
                w.print("Hart {d}: Load page fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}\r\n", .{ hart_id, epc, tval }) catch while (true) {};
                return epc + 4;
            },
            .store_or_amo_page_fault => {
                w.print("Hart {d}: Store/AMO page fault: EPC = 0x{x:0>16}, TVAL = 0x{x:0>16}\r\n", .{ hart_id, epc, tval }) catch while (true) {};
                return epc + 4;
            },
            else => {
                w.print("Hart {d}: Unhandled synchronous interrupt: {d}, EPC = 0x{x:0>16}\r\n", .{ hart_id, cause.num, epc }) catch while (true) {};
                while (true) {}
            },
        }
    }

    return epc;
}

export fn supervisor_trap_vector() 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
        \\
        \\ csrr a0, sepc
        \\ csrr a1, stval
        \\ csrr a2, scause
        // \\ csrr a3, mhartid
        // Use zero for the hart id until a solution is found.
        \\ mv a3, zero
        \\ csrr a4, sstatus
        \\ mv a5, t5
        \\ ld sp, 520(t5)
        \\ call supervisor_trap
        \\
        \\ 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() void {
    trap_frame = .{
        .general_purpose_registers = [_]usize{0} ** 32,
        .floating_point_registers = [_]usize{0} ** 32,
        .satp = 0,
        .stack_pointer = @ptrFromInt(instructions.stackPointer()),
        .hart_id = 0,
    };

    asm volatile (
        \\ csrw sscratch, %[trapframe]
        \\
        \\ la t0, supervisor_trap_vector
        \\ csrw stvec, t0
        :
        : [trapframe] "r" (&trap_frame),
    );
}

pub const setEnabled = instructions.setCsrFn(Enable, "sie").?;