use core::{
    marker::PhantomData,
    ops::{Deref, DerefMut},
};

use alloc::vec::Vec;
use device_api::interrupt::IpiMessage;

use crate::{
    guard::IrqGuard, sync::IrqSafeSpinlock, task::Scheduler, util::OneTimeInit, Architecture,
};

#[repr(C, align(0x10))]
pub struct CpuImpl<A: Architecture, S: Scheduler + 'static> {
    inner: A::PerCpuData,
    scheduler: OneTimeInit<&'static S>,

    id: u32,
    current_thread_id: Option<S::ThreadId>,

    _pd: PhantomData<A>,
}

pub struct LocalCpuImpl<'a, A: Architecture, S: Scheduler + 'static> {
    cpu: &'a mut CpuImpl<A, S>,
    guard: IrqGuard<A>,
}

pub struct IpiQueue<A: Architecture> {
    data: IrqSafeSpinlock<A, Option<IpiMessage>>,
}

pub trait CpuData {
    fn is_bootstrap(&self, id: u32) -> bool {
        // On most architectures
        id == 0
    }

    fn queue_index(&self, id: u32) -> usize {
        // On most architectures
        id as usize
    }
}

pub trait CpuFeatureSet {
    fn iter(&self) -> impl Iterator<Item = &'static str>;
}

impl<A: Architecture, S: Scheduler + 'static> CpuImpl<A, S> {
    pub fn new(id: u32, inner: A::PerCpuData) -> Self {
        Self {
            inner,
            scheduler: OneTimeInit::new(),
            id,
            current_thread_id: None,
            _pd: PhantomData,
        }
    }

    pub fn init_ipi_queues(cpu_count: usize) {
        let queues = Vec::from_iter((0..cpu_count).map(|_| IpiQueue::new()));
        unsafe { A::init_ipi_queues(queues) }
    }

    pub fn is_bootstrap(&self) -> bool {
        self.inner.is_bootstrap(self.id)
    }

    pub fn queue_index(&self) -> usize {
        self.inner.queue_index(self.id)
    }

    pub fn set_current_thread_id(&mut self, id: Option<S::ThreadId>) {
        self.current_thread_id = id;
    }

    pub fn current_thread_id(&self) -> Option<S::ThreadId> {
        self.current_thread_id
    }

    pub fn set_scheduler(&mut self, sched: &'static S) {
        self.scheduler.init(sched);
    }

    pub fn try_get_scheduler(&self) -> Option<&'static S> {
        self.scheduler.try_get().copied()
    }

    pub fn scheduler(&self) -> &'static S {
        self.scheduler.get()
    }

    /// # Safety
    ///
    /// See [Architecture::set_local_cpu].
    pub unsafe fn set_local(&'static mut self) {
        A::set_local_cpu(self as *mut _ as *mut _)
    }

    pub fn try_local<'a>() -> Option<LocalCpuImpl<'a, A, S>> {
        let guard = IrqGuard::acquire();
        let cpu = A::local_cpu() as *mut Self;

        unsafe { cpu.as_mut().map(|cpu| LocalCpuImpl { cpu, guard }) }
    }

    pub fn local<'a>() -> LocalCpuImpl<'a, A, S> {
        Self::try_local().expect("Local CPU not initialized")
    }

    pub fn id(&self) -> u32 {
        self.id
    }

    pub fn push_ipi_queue(cpu_id: u32, msg: IpiMessage) {
        if let Some(queue) = A::ipi_queue(cpu_id) {
            queue.push(msg);
        }
    }

    pub fn get_ipi(&self) -> Option<IpiMessage> {
        A::ipi_queue(self.id).and_then(|q| q.pop())
    }

    pub fn available_features(&self) -> Option<&A::CpuFeatures> {
        A::cpu_available_features(self)
    }

    pub fn enabled_features(&self) -> Option<&A::CpuFeatures> {
        A::cpu_enabled_features(self)
    }
}

impl<A: Architecture, S: Scheduler> Deref for CpuImpl<A, S> {
    type Target = A::PerCpuData;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<A: Architecture, S: Scheduler> DerefMut for CpuImpl<A, S> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

impl<A: Architecture, S: Scheduler + 'static> LocalCpuImpl<'_, A, S> {
    pub fn into_guard(self) -> IrqGuard<A> {
        self.guard
    }
}

impl<A: Architecture, S: Scheduler> Deref for LocalCpuImpl<'_, A, S> {
    type Target = CpuImpl<A, S>;

    fn deref(&self) -> &Self::Target {
        self.cpu
    }
}

impl<A: Architecture, S: Scheduler> DerefMut for LocalCpuImpl<'_, A, S> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.cpu
    }
}

impl<A: Architecture> IpiQueue<A> {
    pub const fn new() -> Self {
        Self {
            data: IrqSafeSpinlock::<A, _>::new(None),
        }
    }

    pub fn push(&self, msg: IpiMessage) {
        let mut lock = self.data.lock();

        assert!(lock.is_none());
        lock.replace(msg);
    }

    pub fn pop(&self) -> Option<IpiMessage> {
        let mut lock = self.data.lock();
        lock.take()
    }
}

impl CpuFeatureSet for () {
    fn iter(&self) -> impl Iterator<Item = &'static str> {
        core::iter::empty()
    }
}
