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

extern crate alloc;

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

use aarch64_cpu::{
    asm::barrier,
    registers::{DAIF, MPIDR_EL1, TPIDR_EL1},
};
use alloc::{boxed::Box, sync::Arc, vec::Vec};
use device_api::interrupt::LocalInterruptController;
use kernel_arch_interface::{
    cpu::{CpuData, CpuImpl, IpiQueue},
    guard::IrqGuard,
    task::Scheduler,
    util::OneTimeInit,
    Architecture,
};
use tock_registers::interfaces::{ReadWriteable, Readable, Writeable};

pub mod context;
pub mod mem;

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

pub struct ArchitectureImpl;

pub trait GicInterface: LocalInterruptController {}

pub struct PerCpuData {
    pub gic: OneTimeInit<Arc<dyn GicInterface>>,
}

impl CpuData for PerCpuData {}

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!("1: nop; b 1b");
}

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

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

    const BREAKPOINT_VALUE: Self::BreakpointType = 0x200020D4;

    fn cpu_index<S: Scheduler + 'static>() -> u32 {
        (MPIDR_EL1.get() & 0xFF) as u32
    }

    fn interrupt_mask() -> bool {
        DAIF.read(DAIF::I) != 0
    }

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

    fn wait_for_interrupt() {
        aarch64_cpu::asm::wfi();
    }

    fn halt() -> ! {
        let _irq = IrqGuard::<Self>::acquire();
        loop {
            aarch64_cpu::asm::wfi();
        }
    }

    unsafe fn set_local_cpu(cpu: *mut ()) {
        TPIDR_EL1.set(cpu as _);
    }

    unsafe fn init_local_cpu<S: Scheduler + 'static>(id: Option<u32>, data: Self::PerCpuData) {
        assert!(
            id.is_none(),
            "AArch64 uses MPIDR_EL1 instead of manual ID set"
        );
        let id = (MPIDR_EL1.get() & 0xFF) as u32;
        let cpu = Box::leak(Box::new(CpuImpl::<Self, S>::new(id, data)));

        cpu.set_local();
    }

    fn local_cpu() -> *mut () {
        TPIDR_EL1.get() as _
    }

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

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

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

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

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

    fn cpu_available_features<S: Scheduler>(_cpu: &CpuImpl<Self, S>) -> Option<&Self::CpuFeatures> {
        None
    }

    fn cpu_enabled_features<S: Scheduler>(_cpu: &CpuImpl<Self, S>) -> Option<&Self::CpuFeatures> {
        None
    }

    // Cache/barrier operation
    fn load_barrier() {
        barrier::dmb(barrier::ISHLD);
    }

    fn store_barrier() {
        barrier::dmb(barrier::ISHST);
    }

    fn memory_barrier() {
        barrier::dsb(barrier::SY);
    }

    fn flush_virtual_range(range: core::ops::Range<usize>) {
        // TODO cache line assumed to be 64 bytes
        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::asm!("dc ivac, {address}", address = in(reg) line);
            }
        }
    }
}
