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

pub var lapic_timer_khz: usize = 0;

// 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 {
    // Get the APIC ready (call first)
    pub fn initialSetup() void {
        singleton.setSpuriousInterruptRegister(.{
            .apic_soft_enable = true,
            .idt_entry = 0xFF,
        });
        // lapic.setTaskPriorityRegister(.{
        //     .priority_class = 0,
        //     .priority_sub_class = 0,
        // });
    }

    pub 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("APIC timer: {} kHz", .{lapic_timer_khz});
    }

    pub fn enablePeriodicInterrupt(ms: usize) void {
        singleton.setInitialCountRegister(0);
        singleton.setDivideConfigurationRegister(.div2);

        singleton.setLVTTimerRegister(.{
            .idt_entry = 48,
            .mode = .periodic,
            .masked = false,
        });

        const ticks: u32 = @truncate(lapic_timer_khz * ms);

        singleton.setInitialCountRegister(ticks);
    }
};

pub fn spurious_interrupt_handler(_: *arch.structures.Idt.InterruptStackFrame) callconv(.{ .x86_64_interrupt = .{} }) void {
    log.warn("Got a spurious interrupt!", .{});
}

pub fn periodic_handler(_: *arch.structures.Idt.InterruptStackFrame) callconv(.{ .x86_64_interrupt = .{} }) void {
    log.warn("Got an ACPI timer interrupt!", .{});
    singleton.setRegister(.eoi, 0);
}