use core::sync::atomic::{AtomicU64, Ordering};
use x86_64::registers::model_specific::Msr;

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

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

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

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

pub fn cpu_id() -> u32 {
    let mut info: u32;
    unsafe { asm!("cpuid", inout("eax") 1 => _, out("ebx") info) }
    if !is_x2apic() {
        info >>= 24;
    }
    info
}

#[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,
}
const NONE_APIC: Option<Apic> = None;

static mut APIC_INIT: [Option<Apic>; CPU_COUNT] = [NONE_APIC; CPU_COUNT];

pub fn move_local_apic(v: u64) -> bool {
    match get_local_apic() {
        None => return false,
        Some(apic) => apic.regs.store(v, Ordering::Release),
    }
    unsafe { Msr::new(APIC_BASE_MSR).write(v) };
    true
}

/// Tries to initialize the local apic at it's default location.
/// returns false if no apic is available
///
/// # Safety
/// Calling this function in a running system might cause the loss of interrupts or other weird
/// behaviour
pub unsafe fn try_initialize_local_apic() -> &'static mut Option<Apic> {
    if is_apic() {
        let cpu_id = cpu_id() as usize;
        if APIC_INIT[cpu_id].is_none() {
            APIC_INIT[cpu_id] = Some(Apic::from_base_apic(APIC_BASE_ADDR));
            Apic::disable_pic();
        }
        APIC_INIT[cpu_id].as_mut().map(|x| x.initialize());
    }
    get_local_apic()
}

pub fn get_local_apic() -> &'static mut Option<Apic> {
    let cpu_id = cpu_id() as usize;
    unsafe { &mut APIC_INIT[cpu_id] }
}

pub struct Apic {
    regs: AtomicU64,
}

impl Apic {
    fn initialize(&mut self) -> Self {
        Self::disable_pic();
        let 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);

        apic
    }

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

    const fn from_base_apic(base_apic: u64) -> Self {
        Self {
            regs: AtomicU64::new((base_apic) & 0xfffffffff << 12),
        }
    }

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

    fn ptr_mut(&mut self) -> *mut u32 {
        self.regs.load(Ordering::SeqCst) 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)
    }
}