#![no_std]
#![allow(clippy::new_without_default)]
#![feature(naked_functions, trait_upcasting)]

extern crate alloc;

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

use alloc::vec::Vec;
use device_api::interrupt::{LocalInterruptController, MessageInterruptController};
use kernel_arch_interface::{
    cpu::{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: &'static dyn LocalApicInterface,
    pub available_features: CpuFeatures,
    pub enabled_features: CpuFeatures,
}

impl PerCpuData {
    pub fn local_apic(&self) -> &'static dyn LocalApicInterface {
        self.local_apic
    }
}

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

#[naked]
extern "C" fn idle_task(_: usize) -> ! {
    unsafe {
        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 local_interrupt_controller() -> &'static dyn LocalInterruptController {
        let local = Self::local_cpu_data().unwrap();
        local.local_apic
    }

    fn message_interrupt_controller() -> &'static dyn MessageInterruptController {
        let local = Self::local_cpu_data().unwrap();
        local.local_apic
    }

    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)
    }
}