use x86_64::registers::model_specific::Msr;

#[cfg(target_arch = "x86_64")]
const APIC_FLAG: u64 = 0x0800;
const APIC_BASE_MSR: u32 = 0x001b;

// APIC constant register values
const APIC_DISABLE: u32 = 0x0001_0000;
const APIC_NMI: u32 = 0x0000_0400;
const APIC_SW_ENABLE: u32 = 0x0000_0100;

pub fn is_x2apic() -> bool {
    let info: u32;
    unsafe {
        asm!("cpuid" : "={ecx}" (info) : "{eax}" (1) : "memory");
    }
    info & (1 << 21) != 0
}

pub fn is_apic() -> bool {
    let info: u32;
    unsafe {
        asm!("cpuid" : "={edx}" (info) : "{eax}" (1) : "memory");
    }
    info & (1 << 9) != 0
}

pub fn set_apic_base(v: u64) {
    unsafe { Msr::new(APIC_BASE_MSR).write(v) }
}

#[repr(C)]
pub struct InterruptCommand {
    low: u32,
    _rsv: [u32; 3],
    high: u32,
}

#[repr(isize)]
pub enum ApicRegister {
    ApicId = 8,
    ApicVer = 12,
    TaskPriority = 32,
    ArbitrationPriority = 36,
    ProcessorPriority = 40,
    EndOfInterrupt = 44,
    LogicalDst = 52,
    DstFmt = 56,
    SpuriousInterruptVec = 60,
    InService = 64,
    TriggerMode = 96,
    InterruptRequest = 128,
    ErrorStatus = 160,
    InterruptCommand = 192,
    InterruptCommandHigh = 196,
    LvtTimer = 200,
    LvtThermalSensor = 204,
    LvtPerformanceMonitor = 208,
    LvtLint0 = 212,
    LvtLint1 = 216,
    LvtError = 220,
    TimerInitialCount = 224,
    TimerCurrentCount = 228,
    TimerDivideConfig = 248,
    ExtApicFeature = 256,
    ExtApicControl = 260,
    SpecificEndOfInterrupt = 264,
    InterruptEnable = 288,
    ExtInterruptLocalVectorTable = 320,
}

#[repr(u32)]
pub enum TimerDivideConfig {
    Div2 = 0,
    Div4 = 1,
    Div8 = 2,
    Div16 = 3,
    Div32 = 8,
    Div64 = 9,
    Div128 = 10,
    NoDivide = 11,
}

static mut LOCAL_APIC: Option<Apic> = None;

pub unsafe fn set_local_apic(apic: Apic) -> &'static mut Apic {
    LOCAL_APIC = Some(apic);
    LOCAL_APIC.as_mut().unwrap()
}

pub unsafe fn get_local_apic() -> Option<&'static mut Apic> {
    LOCAL_APIC.as_mut()
}

pub struct Apic {
    regs: *mut [u32; 1024],
}

impl Apic {
    pub fn new() -> Option<Self> {
        if !is_apic() {
            return None;
        }
        Self::disable_pic();
        let mut base_apic = unsafe { Msr::new(APIC_BASE_MSR).read() };
        let mut apic = Self::from_base_apic(base_apic);

        apic.set(ApicRegister::LogicalDst, 0);

        apic.set(ApicRegister::TaskPriority, 0);
        apic.set(ApicRegister::LvtTimer, APIC_DISABLE);
        apic.set(ApicRegister::LvtPerformanceMonitor, APIC_NMI);
        apic.set(ApicRegister::LvtLint0, APIC_DISABLE);
        apic.set(ApicRegister::LvtLint1, APIC_DISABLE);
        apic.set(ApicRegister::DstFmt, 0x0FFFFFFF);
        apic.set(
            ApicRegister::LvtError,
            super::InterruptType::ApicError.as_u8().into(),
        );
        apic.set(ApicRegister::ErrorStatus, 0);

        unsafe { Msr::new(APIC_BASE_MSR).write(base_apic | APIC_FLAG) };

        apic.set(ApicRegister::SpuriousInterruptVec, 0x1ff);

        Some(apic)
    }

    pub fn disable_pic() {
        unsafe {
            x86_64::instructions::port::PortWriteOnly::new(0xa1).write(0xffu8);
            x86_64::instructions::port::PortWriteOnly::new(0x21).write(0xffu8);
        }
    }

    fn from_base_apic(base_apic: u64) -> Self {
        Self {
            regs: ((base_apic) & 0xfffffffff << 12) as *mut [u32; 1024],
        }
    }

    pub fn ptr(&self) -> *const u32 {
        self.regs as *const u32
    }

    fn ptr_mut(&mut self) -> *mut u32 {
        self.regs as *mut u32
    }

    pub fn get(&self, register: ApicRegister) -> u32 {
        unsafe { self.ptr().offset(register as isize).read_volatile() }
    }

    pub fn set(&mut self, register: ApicRegister, val: u32) {
        unsafe { self.ptr_mut().offset(register as isize).write_volatile(val) }
    }

    pub fn set_interrupt_command(&mut self, dst: u8, v: u8) {
        let (dst, v): (u32, u32) = (dst.into(), v.into());
        self.set(ApicRegister::InterruptCommandHigh, dst << 24);
        self.set(ApicRegister::InterruptCommand, v)
    }

    pub fn end_of_interrupt(&mut self) {
        self.set(ApicRegister::EndOfInterrupt, 0);
    }

    pub fn set_timer_interrupt_handler(
        &mut self,
        divide: TimerDivideConfig,
        intr: super::InterruptType,
    ) {
        self.set(ApicRegister::LvtTimer, intr.as_u8() as u32 | (1 << 17));
        self.set(ApicRegister::TimerDivideConfig, divide as u32);
        self.set(ApicRegister::TimerInitialCount, 1 << 27);
    }

    pub fn get_timer_value(&self) -> u32 {
        self.get(ApicRegister::TimerCurrentCount)
    }

    pub fn get_error_code(&mut self) -> u32 {
        self.set(ApicRegister::ErrorStatus, 0);
        // write contents of apic error register to memory
        self.get(ApicRegister::ErrorStatus)
    }

    pub fn get_id(&self) -> u32 {
        self.get(ApicRegister::ApicId)
    }
}