#![feature(decl_macro)]
#![no_std]

extern crate alloc;

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

use alloc::{boxed::Box, collections::btree_map::BTreeMap, vec::Vec};
use device_api::interrupt::LocalInterruptController;
use kernel_arch_interface::{
    cpu::{CpuData, CpuImpl, IpiQueue},
    sync::IrqSafeSpinlock,
    task::Scheduler,
    util::OneTimeInit,
    Architecture,
};
use tock_registers::interfaces::{ReadWriteable, Readable};

use registers::SSTATUS;

pub mod mem;
pub use mem::{process::ProcessAddressSpaceImpl, KernelTableManagerImpl};
pub mod context;
pub use context::TaskContextImpl;
pub mod intrinsics;
pub mod registers;
pub mod sbi;

pub struct ArchitectureImpl;

#[repr(C)]
pub struct PerCpuData {
    // Used in assembly
    pub tmp_t0: usize,   // 0x00
    pub umode_sp: usize, // 0x08
    pub smode_sp: usize, // 0x10

    // Used elsewhere
    pub bootstrap: bool,
    pub queue_index: usize,
}

pub static CPU_COUNT: AtomicUsize = AtomicUsize::new(1);
pub static mut BOOT_HART_ID: u64 = 0;
static IPI_QUEUES: OneTimeInit<Vec<IpiQueue<ArchitectureImpl>>> = OneTimeInit::new();
static HART_TO_QUEUE: IrqSafeSpinlock<ArchitectureImpl, BTreeMap<u32, usize>> =
    IrqSafeSpinlock::new(BTreeMap::new());

impl CpuData for PerCpuData {
    fn is_bootstrap(&self, id: u32) -> bool {
        let _ = id;
        self.bootstrap
    }

    fn queue_index(&self, id: u32) -> usize {
        let _ = id;
        self.queue_index
    }
}

/// Returns the ID of the bootstrap HART
pub fn boot_hart_id() -> u64 {
    unsafe { BOOT_HART_ID }
}

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

impl ArchitectureImpl {
    pub fn for_each_hart<F: FnMut(u32, usize, &IpiQueue<ArchitectureImpl>)>(mut f: F) {
        let map = HART_TO_QUEUE.lock();
        map.iter().for_each(|(&hart_id, &queue_index)| {
            let queue = &IPI_QUEUES.get()[queue_index];
            f(hart_id, queue_index, queue);
        });
    }
}

impl Architecture for ArchitectureImpl {
    type PerCpuData = PerCpuData;
    type CpuFeatures = ();
    type BreakpointType = u32;

    const BREAKPOINT_VALUE: Self::BreakpointType = 0;

    fn halt() -> ! {
        loop {
            unsafe { Self::set_interrupt_mask(true) };
            Self::wait_for_interrupt();
        }
    }

    unsafe fn set_local_cpu(cpu: *mut ()) {
        unsafe { core::arch::asm!("mv tp, {0}", in(reg) cpu) };
    }

    #[inline]
    fn local_cpu() -> *mut () {
        let value: u64;
        unsafe { core::arch::asm!("mv {0}, tp", out(reg) value) };
        value as _
    }

    unsafe fn init_local_cpu<S: Scheduler + 'static>(id: Option<u32>, data: Self::PerCpuData) {
        let id = id.expect("riscv64 requires an explicit HART ID in its per-processor struct");
        let queue_index = data.queue_index;
        HART_TO_QUEUE.lock().insert(id, queue_index);
        let cpu = Box::leak(Box::new(CpuImpl::<Self, S>::new(id, data)));
        unsafe { cpu.set_local() };
    }

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

    fn ipi_queue(cpu_id: u32) -> Option<&'static IpiQueue<Self>> {
        let queue_index = *HART_TO_QUEUE.lock().get(&cpu_id)?;
        IPI_QUEUES.try_get().and_then(|q| q.get(queue_index))
    }

    #[inline]
    unsafe fn set_interrupt_mask(mask: bool) -> bool {
        let old = Self::interrupt_mask();
        if mask {
            SSTATUS.modify(SSTATUS::SIE::CLEAR);
        } else {
            SSTATUS.modify(SSTATUS::SIE::SET);
        }
        old
    }

    #[inline]
    fn interrupt_mask() -> bool {
        SSTATUS.matches_all(SSTATUS::SIE::CLEAR)
    }

    fn wait_for_interrupt() {
        unsafe {
            core::arch::asm!("wfi");
        }
    }

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

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

    fn cpu_enabled_features<S: Scheduler>(cpu: &CpuImpl<Self, S>) -> Option<&Self::CpuFeatures> {
        let _ = cpu;
        todo!()
    }

    fn cpu_available_features<S: Scheduler>(cpu: &CpuImpl<Self, S>) -> Option<&Self::CpuFeatures> {
        let _ = cpu;
        todo!()
    }

    fn local_interrupt_controller() -> Option<&'static dyn LocalInterruptController> {
        todo!()
    }

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

    // Cache/barrier operation
    fn load_barrier() {
        unsafe { core::arch::asm!("fence r, w") };
    }

    fn store_barrier() {
        unsafe { core::arch::asm!("fence w, r") };
    }

    fn memory_barrier() {
        unsafe { core::arch::asm!("fence rw, rw") };
    }

    fn flush_virtual_range(_range: Range<usize>) {
        // TODO
    }
}
