// 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 TrapFrame = @import("TrapFrame.zig");
const plic = @import("plic.zig");
const process = @import("process.zig");
const riscv = @import("riscv.zig");
const syscall = @import("syscall.zig");

pub var trap_frame: TrapFrame = undefined;

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

    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,
    supervisor_software,
    user_timer = 4,
    supervisor_timer,
    user_external = 8,
    supervisor_external,
    _,
};

pub const SyncCause = enum(u63) {
    instruction_address_misaligned,
    instruction_access_fault,
    illegal_instruction,
    breakpoint,
    load_access_fault = 5,
    amo_address_misaligned,
    store_or_amo_access_fault,
    ecall,
    instruction_page_fault = 12,
    load_page_fault,
    store_or_amo_page_fault = 15,
    _,
};

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

    const status = riscv.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 => {
                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("schedule error: {}", .{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("plic claim error: {}", .{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("plic complete error: {}", .{err});
                    };
                }
            },
            else => {
                std.debug.panic("unhandled interrupt {}", .{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("illegal instruction", .{});
            },
            .instruction_access_fault => {
                std.debug.panic("instruction access fault", .{});
            },
            .load_access_fault => {
                std.debug.panic("load access fault", .{});
            },
            .store_or_amo_access_fault => {
                std.debug.panic("store/amo access fault", .{});
            },
            .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("instruction page fault", .{});
            },
            .load_page_fault => {
                std.debug.panic("load page fault", .{});
            },
            .store_or_amo_page_fault => {
                std.debug.panic("store/amo page fault", .{});
            },
            else => {
                std.debug.panic("unhandled exception {d}", .{cause.num});
            },
        }
    }

    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, scause
        \\ mv a2, 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(riscv.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),
    );
}