const std = @import("std");
const arch = @import("../root.zig");
const idt = arch.interrupts.idt;
const log = std.log.scoped(.apic);
const common = @import("common");

pub var lapic_timer_khz: usize = 0;
pub var tsc_deadline_available = false;
pub var has_x2apic: bool = false;

// tbh every cpu will be either x2apic or not, and if xapic it will
// have the exact same base address anyways so this is fine
pub var singleton: LAPIC = undefined;

// Must instantiate this!
pub const LAPIC = union(enum) {
    xapic: [*]volatile u8,
    x2apic,

    const Self = @This();

    // ID Register

    pub const IdRegister = packed struct(u32) {
        _reserved0: u24 = 0,
        id: u8,
    };

    pub fn getIdRegister(lapic: Self) IdRegister {
        return @bitCast(lapic.getRegister(.lapic_id));
    }

    // Version Register
    pub const VersionRegister = packed struct(u32) {
        version: u8,
        _reserved0: u8 = 0,
        max_lvt_entry: u8,
        support_for_eoi_broadcast_suppression: bool,
        _reserved1: u7 = 0,
    };

    pub fn getVersionRegister(lapic: Self) VersionRegister {
        return @bitCast(lapic.getRegister(.version));
    }

    // Spurious Interrupt
    pub const SpuriousInterruptRegister = packed struct(u32) {
        idt_entry: u8,
        apic_soft_enable: bool,
        focus_processor_checking: bool = false,
        _reserved0: u2 = 0,
        eoi_broadcast_suppression: bool = false,
        _reserved1: u19 = 0,
    };

    pub fn getSpuriousInterruptRegister(lapic: Self) SpuriousInterruptRegister {
        return @bitCast(lapic.getRegister(.spurious_vector));
    }

    pub fn setSpuriousInterruptRegister(lapic: Self, val: SpuriousInterruptRegister) void {
        lapic.setRegister(.spurious_vector, @bitCast(val));
    }

    // Task Priority
    pub const TaskPriorityRegister = packed struct(u32) {
        priority_sub_class: u4,
        priority_class: u4,
        _reserved0: u24 = 0,
    };

    pub fn getTaskPriorityRegister(lapic: Self) TaskPriorityRegister {
        return @bitCast(lapic.getRegister(.task_priority));
    }

    pub fn setTaskPriorityRegister(lapic: Self, val: TaskPriorityRegister) void {
        lapic.setRegister(.task_priority, @bitCast(val));
    }

    // Initial Count
    pub fn setInitialCountRegister(lapic: Self, count: u32) void {
        lapic.setRegister(.initial_count, count);
    }

    // Current Count
    pub fn getCurrentCountRegister(lapic: Self) u32 {
        return lapic.getRegister(.current_count);
    }

    // Divide Configuration
    pub const DivideConfigurationRegister = enum(u32) {
        div2 = 0,
        div4 = 1,
        div8 = 2,
        div16 = 3,
        div32 = 8,
        div64 = 9,
        div128 = 10,
        div1 = 11,
    };

    pub fn getDivideConfigurationRegister(lapic: LAPIC) DivideConfigurationRegister {
        return @enumFromInt(lapic.getRegister(.divide_configuration));
    }

    pub fn setDivideConfigurationRegister(lapic: LAPIC, register: DivideConfigurationRegister) void {
        lapic.setRegister(.divide_configuration, @intFromEnum(register));
    }

    // LVT Timer Register
    pub const LVTTimerRegister = packed struct(u32) {
        idt_entry: u8,
        _reserved0: u4 = 0,
        status: DeliveryStatus = .idle,
        _reserved1: u3 = 0,
        masked: bool,
        mode: Mode,
        _reserved2: u13 = 0,

        pub const DeliveryStatus = enum(u1) {
            idle = 0,
            send_pending = 1,
        };

        pub const Mode = enum(u2) {
            oneshot = 0,
            periodic = 1,
            tsc_deadline = 2,
        };
    };

    pub fn getLVTTimerRegister(lapic: LAPIC) LVTTimerRegister {
        return @enumFromInt(lapic.getRegister(.lvt_timer));
    }

    pub fn setLVTTimerRegister(lapic: LAPIC, register: LVTTimerRegister) void {
        lapic.setRegister(.lvt_timer, @bitCast(register));
    }

    pub fn getRegister(lapic: Self, reg: Register) u32 {
        switch (lapic) {
            .xapic => |base| {
                const ptr: *align(0x10) volatile u32 = @ptrCast(@alignCast(base + reg.xapic()));
                return ptr.*;
            },
            .x2apic => {
                return arch.registers.readMSR(u32, reg.x2apic());
            },
        }
    }

    pub fn setRegister(lapic: Self, reg: Register, value: u32) void {
        switch (lapic) {
            .xapic => |base| {
                const ptr: *align(0x10) volatile u32 = @ptrCast(@alignCast(base + reg.xapic()));
                ptr.* = value;
            },
            .x2apic => {
                arch.registers.writeMSR(u32, reg.x2apic(), value);
            },
        }
    }

    pub const Register = enum(u32) {
        // From Intel® 64 and IA-32 Architectures Software Developer’s Manual Volume 3
        lapic_id = 0x2,
        version = 0x3,
        task_priority = 0x8,
        process_priority = 0xa,
        eoi = 0xb,
        logical_destination = 0xd,
        spurious_vector = 0xf,

        in_service_0_31 = 0x10,
        in_service_63_32 = 0x11,
        in_service_95_64 = 0x12,
        in_service_127_96 = 0x13,
        in_service_159_128 = 0x14,
        in_service_191_160 = 0x15,
        in_service_223_192 = 0x16,
        in_service_255_224 = 0x17,

        trigger_mode_0_31 = 0x18,
        trigger_mode_63_32 = 0x19,
        trigger_mode_95_64 = 0x1a,
        trigger_mode_127_96 = 0x1b,
        trigger_mode_159_128 = 0x1c,
        trigger_mode_191_160 = 0x1d,
        trigger_mode_223_192 = 0x1e,
        trigger_mode_255_224 = 0x1f,

        interrupt_request_0_31 = 0x20,
        interrupt_request_63_32 = 0x21,
        interrupt_request_95_64 = 0x22,
        interrupt_request_127_96 = 0x23,
        interrupt_request_159_128 = 0x24,
        interrupt_request_191_160 = 0x25,
        interrupt_request_223_192 = 0x26,
        interrupt_request_255_224 = 0x27,

        error_status = 0x28,
        lvt_cmi = 0x2f,

        interrupt_command_0_31 = 0x30,
        interrupt_command_32_63 = 0x31,
        lvt_timer = 0x32,
        lvt_thermal_sensor = 0x33,
        lvt_performance_monitoring = 0x34,
        lvt_lint0 = 0x35,
        lvt_lint1 = 0x36,
        lvt_error = 0x37,
        initial_count = 0x38,
        current_count = 0x39,
        divide_configuration = 0x3e,
        self_ipi = 0x3f,

        const Self = @This();

        // Get an offset to apply to the xAPIC base
        pub fn xapic(reg: Register) usize {
            return @intFromEnum(reg) * 0x10;
        }

        // Get an MSR number to write to
        pub fn x2apic(reg: Register) u32 {
            return 0x800 | @intFromEnum(reg);
        }
    };
};

pub const init = struct {
    pub fn initialSetup() void {
        // First, make the APIC accessible
        initSingleton() catch |err| {
            log.err("Failed to map APIC! {}", .{err});
            @panic("initSingleton");
        };
        // Set up the interrupt handlers
        singleton.setSpuriousInterruptRegister(.{
            .apic_soft_enable = true,
            .idt_entry = 0xFF,
        });
        // lapic.setTaskPriorityRegister(.{
        //     .priority_class = 0,
        //     .priority_sub_class = 0,
        // });
        arch.interrupts.idt.add_handler(.{ .interrupt = 0xFF }, u64, spurious_interrupt_handler, 0, 0);
        arch.interrupts.idt.add_handler(.{ .interrupt = 48 }, u64, timer_handler, 0, 0);

        // Calibrate against the TSC
        calibrateTimer();
        // Set up the LVT Timer Register
        enableOneshotInterrupt();
    }

    fn initSingleton() !void {
        arch.interrupts.apic.singleton = switch (has_x2apic) {
            true => .x2apic,
            false => blk: {
                // Map the APIC first!
                const apic_base = common.mm.physToHHDM([*]volatile u8, 0xFEE0_0000);
                try common.mm.paging.mapPhys(.{
                    .vaddr = @intFromPtr(apic_base),
                    .paddr = 0xFEE0_0000,
                    .size = 0x1000,
                    .memory_type = .DeviceUncacheable,
                    .perms = .{
                        .x = false,
                        .u = false,
                        .w = true,
                    },
                });
                break :blk .{ .xapic = apic_base };
            },
        };
    }

    fn calibrateTimer() void {
        singleton.setDivideConfigurationRegister(.div2);
        singleton.setLVTTimerRegister(.{
            .idt_entry = 0x69,
            .mode = .oneshot,
            .masked = true,
        });

        var warmup_clk_count: u64 = 0;
        for (0..5) |_| {
            singleton.setInitialCountRegister(0xFFFF_FFFF);
            arch.tsc.delay_poll(1);
            const count = singleton.getCurrentCountRegister();
            singleton.setInitialCountRegister(0);

            warmup_clk_count += 0xFFFF_FFFF - count;
        }

        std.mem.doNotOptimizeAway(&warmup_clk_count);

        // Now, do it again
        var tick_count: u64 = 0;
        for (0..5) |_| {
            singleton.setInitialCountRegister(0xFFFF_FFFF);
            arch.tsc.delay_poll(5);
            const count = singleton.getCurrentCountRegister();
            singleton.setInitialCountRegister(0);

            tick_count += 0xFFFF_FFFF - count;
        }

        // This would be the number of ticks per 5ms interval
        const norm = tick_count / 5;

        lapic_timer_khz = norm / 5;

        log.debug("timer: {} kHz", .{lapic_timer_khz});
    }

    fn enableOneshotInterrupt() void {
        const mode: LAPIC.LVTTimerRegister.Mode = switch (tsc_deadline_available) {
            true => .tsc_deadline,
            false => blk: {
                singleton.setInitialCountRegister(0);
                singleton.setDivideConfigurationRegister(.div2);
                break :blk .oneshot;
            },
        };

        // TODO: detect and support tsc_deadline, ditto @ armTimer
        singleton.setLVTTimerRegister(.{
            .idt_entry = 48,
            .mode = mode,
            .masked = false,
        });
    }
};

pub fn armTimer(ms: usize) void {
    if (tsc_deadline_available) {
        const IA32_TSC_DEADLINE = arch.registers.MSR(u64, 0x6E0);
        const delta = arch.tsc.tsc_khz * ms;
        const target = arch.tsc.rdtsc() + delta;

        IA32_TSC_DEADLINE.write(target);
    } else {
        const ticks: u32 = @truncate(lapic_timer_khz * ms);
        singleton.setInitialCountRegister(ticks);
    }
}

pub fn spurious_interrupt_handler(_: *idt.InterruptFrame(u64)) callconv(idt.CallConv) void {
    log.warn("Got a spurious interrupt!", .{});
}

pub fn timer_handler(stack_trace: *idt.InterruptFrame(u64)) callconv(idt.CallConv) void {
    defer {
        singleton.setRegister(.eoi, 0);
        armTimer(20);
    }
    // 1. Get the next task. If there is no next task, just keep scheduling.
    const task = common.scheduler.getNextTask() orelse return;
    // 2. Swap the next task state with the current interrupt trace
    std.mem.swap(arch.interrupts.idt.SavedRegisters, &task.regs, &stack_trace.regs);
    std.mem.swap(u64, &task.rip, &stack_trace.rip);
    std.mem.swap(u64, &task.rsp, &stack_trace.rsp);
    // If task has a new cr3, swap current CR3 and task cr3 too
    if (task.cr3_val != stack_trace.cr3) {
        arch.registers.ControlRegisters.Cr3.write(task.cr3_val);
        task.cr3_val = stack_trace.cr3;
    }
    // 3. Now, `task` has our current state, so enqueue it.
    common.scheduler.pushTask(task);
}