//! x86-64 Local APIC driver implementation
use core::sync::atomic::{AtomicUsize, Ordering};

use abi::error::Error;
use alloc::{sync::Arc, vec, vec::Vec};
use device_api::{
    device::Device,
    interrupt::{
        InterruptAffinity, InterruptHandler, IpiDeliveryTarget, IpiMessage, IrqVector,
        LocalInterruptController, MessageInterruptController, MsiInfo,
    },
};
use kernel_arch_x86::registers::MSR_IA32_APIC_BASE;
use kernel_arch_x86_64::{mem::table::L3, LocalApicInterface, CPU_COUNT};
use libk::arch::Cpu;
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo, table::EntryLevelExt};
use libk_util::{
    sync::{spin_rwlock::IrqSafeRwLock, IrqGuard},
    OneTimeInit,
};
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::{ReadOnly, ReadWrite, WriteOnly},
};

use crate::arch::x86_64::apic::APIC_MSI_OFFSET;

use super::{
    APIC_IPI_VECTOR, APIC_LINT0_VECTOR, APIC_LINT1_VECTOR, APIC_SPURIOUS_VECTOR, APIC_TIMER_VECTOR,
    MAX_MSI_VECTORS,
};

const TIMER_INTERVAL: u32 = 15000;

/// When initialized, contains the Local APIC ID of the bootstrap processor
pub static BSP_APIC_ID: OneTimeInit<u32> = OneTimeInit::new();

register_bitfields! {
    u32,
    Id [
        ApicId OFFSET(24) NUMBITS(8) []
    ],
    SpuriousVector [
        Vector OFFSET(0) NUMBITS(8) [],
        SoftwareEnable OFFSET(8) NUMBITS(1) [],
    ],
    TimerLocalVectorEntry [
        Vector OFFSET(0) NUMBITS(8) [],
        Mask OFFSET(16) NUMBITS(1) [
            Masked = 1,
            Unmasked = 0
        ],
        TimerMode OFFSET(17) NUMBITS(1) [
            Periodic = 1,
            OneShot = 0
        ]
    ],
    LocalVectorEntry [
        Vector OFFSET(0) NUMBITS(8) [],
        Mask OFFSET(16) NUMBITS(1) [
            Masked = 1,
            Unmasked = 0,
        ],
        DeliveryMode OFFSET(8) NUMBITS(3) [
            Nmi = 4,
            ExtINT = 7
        ],
    ],
    ICR0 [
        Vector OFFSET(0) NUMBITS(8) [],
        Destination OFFSET(8) NUMBITS(3) [
            Normal = 1,
            Lowest = 2,
            SMI = 3,
            NMI = 4,
            INIT = 5,
            SIPI = 6
        ],
        DeliveryStatus OFFSET(12) NUMBITS(1) [],
        INIT0 OFFSET(14) NUMBITS(1) [
            Deassert = 0,
            Assert = 1,
        ],
        INIT1 OFFSET(15) NUMBITS(1) [
            Deassert = 1,
            Assert = 0,
        ],
        DestinationType OFFSET(18) NUMBITS(3) [
            Physical = 0,
            This = 1,
            All = 2,
            AllExceptThis = 3,
        ]
    ],
    ICR1 [
        PhysicalDestination OFFSET(24) NUMBITS(4) []
    ],
}

register_structs! {
    #[allow(non_snake_case, missing_docs)]
    Regs {
        (0x00 => _0),
        (0x20 => Id: ReadOnly<u32, Id::Register>),
        (0x24 => _1),
        (0x80 => TaskPriorityRegister: ReadWrite<u32>),
        (0x84 => _13),
        (0xB0 => EndOfInterrupt: WriteOnly<u32>),
        (0xB4 => _2),
        (0xF0 => SpuriousVector: ReadWrite<u32, SpuriousVector::Register>),
        (0xF4 => _3),
        (0x100 => ISR0: ReadOnly<u32>),
        (0x104 => _14),
        (0x280 => ErrorStatus: ReadOnly<u32>),
        (0x284 => _4),
        (0x300 => ICR0: ReadWrite<u32, ICR0::Register>),
        (0x304 => _5),
        (0x310 => ICR1: ReadWrite<u32, ICR1::Register>),
        (0x314 => _6),
        (0x320 => TimerLocalVectorEntry: ReadWrite<u32, TimerLocalVectorEntry::Register>),
        (0x324 => _7),
        (0x350 => LInt0: ReadWrite<u32, LocalVectorEntry::Register>),
        (0x354 => _8),
        (0x360 => LInt1: ReadWrite<u32, LocalVectorEntry::Register>),
        (0x364 => _9),
        (0x380 => TimerInitCount: ReadWrite<u32>),
        (0x384 => _10),
        (0x390 => TimerCurrentCount: ReadOnly<u32>),
        (0x394 => _11),
        (0x3E0 => TimerDivideConfig: ReadWrite<u32>),
        (0x3E4 => _12),
        (0x530 => @END),
    }
}

#[derive(Clone)]
struct MsiVector {
    logical: usize,
    handler: Arc<dyn InterruptHandler>,
}

struct MsiVectorTable {
    // MSI_MAX_VECTORS * vectors per row
    msi_vectors: IrqSafeRwLock<Vec<Vec<MsiVector>>>,
    last_msi_allocated: AtomicUsize,
}

impl MsiVectorTable {
    pub fn new(physical_vectors: usize) -> Self {
        let rows = vec![Vec::new(); physical_vectors];

        Self {
            msi_vectors: IrqSafeRwLock::new(rows),
            last_msi_allocated: AtomicUsize::new(0),
        }
    }

    pub fn map(
        &self,
        apic_id: usize,
        range: &mut [MsiInfo],
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<(), Error> {
        if range.is_empty() {
            return Err(Error::InvalidArgument);
        }

        let seq = self
            .last_msi_allocated
            .fetch_add(range.len(), Ordering::AcqRel);
        let mut rows = self.msi_vectors.write();

        let physical_start = seq % rows.len();
        let physical_end = seq.wrapping_add(range.len()) % rows.len();
        if range.len() == 1 {
            log::info!("Map MSI {:?} -> {}", handler.display_name(), physical_start);
        } else {
            log::info!(
                "Map MSI {:?} -> {:?}",
                handler.display_name(),
                physical_start..physical_end - 1
            );
        }

        for (logical, info) in range.iter_mut().enumerate() {
            let physical = logical.wrapping_add(seq) % rows.len();
            rows[physical].push(MsiVector {
                logical,
                handler: handler.clone(),
            });

            let address = 0xFEE00000 | (apic_id << 12);
            let value = 32 + APIC_MSI_OFFSET + physical as u32;

            *info = MsiInfo {
                vector: logical,
                affinity: InterruptAffinity::Specific(0),
                address,
                value,
            };
        }

        Ok(())
    }

    pub fn invoke(&self, physical_vector: usize) {
        let rows = self.msi_vectors.read();
        let row = &rows[physical_vector];

        for col in row.iter() {
            col.handler.clone().handle_irq(IrqVector::Msi(col.logical));
        }
    }
}

/// Per-processor local APIC interface
pub struct LocalApic {
    regs: DeviceMemoryIo<'static, Regs>,
    id: u32,
    msi_table: MsiVectorTable,
}

unsafe impl Send for LocalApic {}
unsafe impl Sync for LocalApic {}

impl Device for LocalApic {
    fn display_name(&self) -> &str {
        "Local APIC"
    }
}

impl LocalApicInterface for LocalApic {
    #[inline]
    fn clear_interrupt(&self) {
        self.regs.EndOfInterrupt.set(0);
    }

    unsafe fn wakeup_cpu(&self, apic_id: u32, entry_vector: PhysicalAddress) {
        log::info!("Waking up apic{}, entry = {:#x}", apic_id, entry_vector);

        while self.regs.ICR0.matches_all(ICR0::DeliveryStatus::SET) {
            core::hint::spin_loop();
        }

        let entry_vector = entry_vector.page_index::<L3>();

        // INIT assert
        self.regs.ICR1.write(ICR1::PhysicalDestination.val(apic_id));

        self.regs.ICR0.write(
            ICR0::Destination::INIT
                + ICR0::DestinationType::Physical
                + ICR0::INIT0::Assert
                + ICR0::INIT1::Assert,
        );

        while self.regs.ICR0.matches_all(ICR0::DeliveryStatus::SET) {
            core::hint::spin_loop();
        }

        // INIT deassert
        self.regs.ICR1.write(ICR1::PhysicalDestination.val(apic_id));

        self.regs.ICR0.write(
            ICR0::Destination::INIT
                + ICR0::DestinationType::Physical
                + ICR0::INIT0::Deassert
                + ICR0::INIT1::Deassert,
        );

        while self.regs.ICR0.matches_all(ICR0::DeliveryStatus::SET) {
            core::hint::spin_loop();
        }

        // Send another SIPI type IPI because the spec says so
        self.regs.ICR1.write(ICR1::PhysicalDestination.val(apic_id));

        self.regs.ICR0.write(
            ICR0::Vector.val(entry_vector as u32)
                + ICR0::Destination::SIPI
                + ICR0::DestinationType::Physical,
        );

        while self.regs.ICR0.matches_all(ICR0::DeliveryStatus::SET) {
            core::hint::spin_loop();
        }
    }
}

impl MessageInterruptController for LocalApic {
    fn handle_msi(&self, vector: usize) {
        self.msi_table.invoke(vector);
    }

    fn register_msi_range(
        self: Arc<Self>,
        range: &mut [MsiInfo],
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<(), Error> {
        let _guard = IrqGuard::acquire();
        self.msi_table.map(self.id as usize, range, handler)
    }
}

impl LocalInterruptController for LocalApic {
    fn send_ipi(&self, target: IpiDeliveryTarget, msg: IpiMessage) -> Result<(), Error> {
        while self.regs.ICR0.matches_all(ICR0::DeliveryStatus::SET) {
            core::hint::spin_loop();
        }

        // TODO use NMI or regular interrupt depending on severity of the message
        match target {
            IpiDeliveryTarget::OtherCpus => {
                let local = Cpu::local();
                let local_id = local.id() as usize;

                for i in 0..CPU_COUNT.load(Ordering::Acquire) {
                    if i != local_id {
                        Cpu::push_ipi_queue(i as u32, msg);
                    }
                }

                self.regs.ICR1.write(ICR1::PhysicalDestination.val(0));
                self.regs.ICR0.write(
                    ICR0::Vector.val(APIC_IPI_VECTOR + 32)
                        + ICR0::Destination::NMI
                        + ICR0::DestinationType::AllExceptThis,
                );

                Ok(())
            }
            IpiDeliveryTarget::ThisCpu => unimplemented!(),
            IpiDeliveryTarget::Specific(_) => unimplemented!(),
        }
    }

    unsafe fn init_ap(&self) -> Result<(), Error> {
        unreachable!()
    }
}

impl LocalApic {
    /// Constructs a new instance of Local APIC.
    ///
    /// # Safety
    ///
    /// Only meant to be called once per processor during their init.
    pub unsafe fn new() -> Self {
        let regs = DeviceMemoryIo::<Regs>::map(Self::base(), Default::default()).unwrap();

        let id = regs.Id.read(Id::ApicId);

        if Self::is_bootstrap_cpu() {
            BSP_APIC_ID.init(id);
        }

        Self::enable();
        // Configure spurious interrupt handler
        regs.SpuriousVector.write(
            SpuriousVector::SoftwareEnable::SET
                + SpuriousVector::Vector.val(APIC_SPURIOUS_VECTOR + 32),
        );

        // Configure task priority register
        regs.TaskPriorityRegister.set(0);

        // Enable timer
        regs.TimerDivideConfig.set(0x2);
        regs.TimerInitCount.set(TIMER_INTERVAL);

        // Configure local interrupt vectors
        regs.TimerLocalVectorEntry.write(
            TimerLocalVectorEntry::Vector.val(APIC_TIMER_VECTOR + 32)
                + TimerLocalVectorEntry::Mask::Unmasked
                + TimerLocalVectorEntry::TimerMode::Periodic,
        );
        // LINT0 unmasked, leave LINT1 masked
        regs.LInt0.write(
            LocalVectorEntry::Mask::Unmasked
                + LocalVectorEntry::Vector.val(APIC_LINT0_VECTOR + 32)
                + LocalVectorEntry::DeliveryMode::ExtINT,
        );
        regs.LInt1.write(
            LocalVectorEntry::Mask::Masked + LocalVectorEntry::Vector.val(APIC_LINT1_VECTOR + 32),
        );

        let msi_table = MsiVectorTable::new(MAX_MSI_VECTORS as _);

        Self {
            id,
            regs,
            msi_table,
        }
    }

    #[inline]
    fn base() -> PhysicalAddress {
        PhysicalAddress::from_u64(MSR_IA32_APIC_BASE.read_base())
    }

    #[inline]
    fn is_bootstrap_cpu() -> bool {
        MSR_IA32_APIC_BASE.read(MSR_IA32_APIC_BASE::BootstrapCpuCore) != 0
    }

    #[inline]
    fn enable() {
        MSR_IA32_APIC_BASE.modify(
            MSR_IA32_APIC_BASE::ApicEnable::SET + MSR_IA32_APIC_BASE::ExtendedEnable::CLEAR,
        );
    }
}