#![no_std]
#![allow(clippy::new_without_default)]

extern crate alloc;

use core::{
    ops::{DerefMut, Range},
    sync::atomic::{AtomicUsize, Ordering},
};

use alloc::{sync::Arc, vec::Vec};
use device_api::interrupt::{LocalInterruptController, MessageInterruptController};
use kernel_arch_interface::{
    cpu::{CpuData, CpuImpl, IpiQueue},
    task::Scheduler,
    util::OneTimeInit,
    Architecture,
};
use kernel_arch_x86::{cpuid::CpuFeatures, registers::MSR_IA32_KERNEL_GS_BASE};
use libk_mm_interface::address::PhysicalAddress;
use tock_registers::interfaces::Writeable;

pub mod context;
pub mod mem;

pub use context::TaskContextImpl;
pub use mem::{process::ProcessAddressSpaceImpl, KernelTableManagerImpl};

pub struct ArchitectureImpl;

pub const KERNEL_VIRT_OFFSET: usize = 0xFFFFFF8000000000;

pub trait LocalApicInterface: LocalInterruptController + MessageInterruptController {
    /// Performs an application processor startup sequence.
    ///
    /// # Safety
    ///
    /// Unsafe: only meant to be called by the BSP during SMP init.
    unsafe fn wakeup_cpu(&self, apic_id: u32, bootstrap_code: PhysicalAddress);

    /// Signals local APIC that we've handled the IRQ
    fn clear_interrupt(&self);
}

#[repr(C, align(0x10))]
pub struct PerCpuData {
    // 0x00
    pub this: *mut Self,
    // 0x08, used in assembly
    pub tss_address: usize,
    // 0x10, used in assembly
    pub tmp_address: usize,

    pub local_apic: Arc<dyn LocalApicInterface>,
    pub available_features: CpuFeatures,
    pub enabled_features: CpuFeatures,
}

impl CpuData for PerCpuData {}

impl PerCpuData {
    #[inline]
    pub fn local_apic(&self) -> &dyn LocalApicInterface {
        self.local_apic.as_ref()
    }
}

static IPI_QUEUES: OneTimeInit<Vec<IpiQueue<ArchitectureImpl>>> = OneTimeInit::new();
pub static CPU_COUNT: AtomicUsize = AtomicUsize::new(1);

#[unsafe(naked)]
extern "C" fn idle_task(_: usize) -> ! {
    core::arch::naked_asm!(
        r#"
            1:
                nop
                jmp 1b
            "#,
        options(att_syntax)
    );
}

impl ArchitectureImpl {
    fn local_cpu_data() -> Option<&'static mut PerCpuData> {
        unsafe { (Self::local_cpu() as *mut PerCpuData).as_mut() }
    }

    fn set_local_tss_sp0(sp: usize) {
        let local_cpu = Self::local_cpu_data().unwrap();
        unsafe {
            (core::ptr::with_exposed_provenance_mut::<usize>(local_cpu.tss_address + 4))
                .write_unaligned(sp);
        }
    }
}

impl Architecture for ArchitectureImpl {
    type PerCpuData = PerCpuData;
    type CpuFeatures = CpuFeatures;
    type BreakpointType = u8;

    const BREAKPOINT_VALUE: Self::BreakpointType = 0xCC;

    unsafe fn set_local_cpu(cpu: *mut ()) {
        MSR_IA32_KERNEL_GS_BASE.set(cpu as u64);
        core::arch::asm!("wbinvd; swapgs");
    }

    fn local_cpu() -> *mut () {
        let mut addr: u64;
        unsafe {
            core::arch::asm!("movq %gs:(0), {0}", out(reg) addr, options(att_syntax));
        }
        addr as _
    }

    unsafe fn init_ipi_queues(queues: Vec<IpiQueue<Self>>) {
        IPI_QUEUES.init(queues);
    }

    unsafe fn init_local_cpu<S: Scheduler + 'static>(id: Option<u32>, data: Self::PerCpuData) {
        use alloc::boxed::Box;

        let cpu = Box::leak(Box::new(CpuImpl::<Self, S>::new(
            id.expect("x86_64 required manual CPU ID set"),
            data,
        )));
        cpu.this = cpu.deref_mut();

        cpu.set_local();
    }

    fn idle_task() -> extern "C" fn(usize) -> ! {
        idle_task
    }

    fn cpu_count() -> usize {
        CPU_COUNT.load(Ordering::Acquire)
    }

    fn cpu_index<S: Scheduler + 'static>() -> u32 {
        CpuImpl::<Self, S>::local().id()
    }

    fn interrupt_mask() -> bool {
        let mut flags: u64;
        unsafe {
            core::arch::asm!("pushfq; pop {0}", out(reg) flags, options(att_syntax));
        }
        // If IF is zero, interrupts are disabled (masked)
        flags & (1 << 9) == 0
    }

    unsafe fn set_interrupt_mask(mask: bool) -> bool {
        let old = Self::interrupt_mask();
        if mask {
            core::arch::asm!("cli");
        } else {
            core::arch::asm!("sti");
        }
        old
    }

    #[inline]
    fn wait_for_interrupt() {
        unsafe {
            core::arch::asm!("hlt");
        }
    }

    fn halt() -> ! {
        loop {
            unsafe {
                core::arch::asm!("cli; hlt");
            }
        }
    }

    fn ipi_queue(cpu_id: u32) -> Option<&'static IpiQueue<Self>> {
        IPI_QUEUES.try_get().and_then(|q| q.get(cpu_id as usize))
    }

    fn local_interrupt_controller() -> Option<&'static dyn LocalInterruptController> {
        let cpu = Self::local_cpu_data()?;
        Some(cpu.local_apic.as_ref())
    }

    fn cpu_enabled_features<S: Scheduler>(cpu: &CpuImpl<Self, S>) -> Option<&Self::CpuFeatures> {
        Some(&cpu.enabled_features)
    }

    fn cpu_available_features<S: Scheduler>(cpu: &CpuImpl<Self, S>) -> Option<&Self::CpuFeatures> {
        Some(&cpu.available_features)
    }

    // Cache/barrier

    fn load_barrier() {
        unsafe { core::arch::x86_64::_mm_lfence() };
    }

    fn store_barrier() {
        unsafe { core::arch::x86_64::_mm_sfence() };
    }

    fn memory_barrier() {
        unsafe { core::arch::x86_64::_mm_mfence() };
    }

    fn flush_virtual_range(range: Range<usize>) {
        // TODO I assume 64-byte cache line on all CPUs
        // TODO clflush instruction may not be available, test for it
        const CLSIZE: usize = 64;
        let start = range.start & !(CLSIZE - 1);
        let end = (range.end + (CLSIZE - 1)) & !(CLSIZE - 1);

        for line in (start..end).step_by(CLSIZE) {
            unsafe { core::arch::x86_64::_mm_clflush(line as _) };
        }
    }
}
