use core::{
    cell::Cell,
    sync::atomic::{AtomicU64, Ordering},
    time::Duration,
};

use alloc::{sync::Arc, vec::Vec};
use crossbeam_queue::SegQueue;
use kernel_arch::{
    task::{Scheduler, TaskContext},
    Architecture, ArchitectureImpl, CpuImpl,
};
use libk_util::{sync::IrqGuard, OneTimeInit};
use yggdrasil_abi::time::SystemTime;

use crate::{
    task::{thread::Thread, TaskContextImpl, ThreadId, ThreadState},
    time::monotonic_time,
};

/// Per-CPU queue
pub struct CpuQueue {
    queue: SegQueue<ThreadId>,
    index: usize,
    idle: Cell<TaskContextImpl>,

    stats: SchedulerStats,
    last_stats_measure: Cell<Option<SystemTime>>,
}

// TODO sort out Cell<...>
unsafe impl Sync for CpuQueue {}

#[derive(Default)]
pub struct SchedulerStats {
    pub idle: AtomicU64,
    pub user: AtomicU64,
    pub kernel: AtomicU64,
    pub total: AtomicU64,
}

static QUEUES: OneTimeInit<Vec<CpuQueue>> = OneTimeInit::new();

impl CpuQueue {
    /// Creates a new [CpuQueue] for CPU with given `index`
    pub fn new(index: usize) -> Self {
        let idle = TaskContextImpl::kernel(
            ArchitectureImpl::idle_task(),
            CpuImpl::<Self>::local().queue_index(),
        )
        .expect("Could not construct an idle task");

        Self {
            queue: SegQueue::new(),
            index,
            idle: Cell::new(idle),

            stats: SchedulerStats::default(),
            last_stats_measure: Cell::new(None),
        }
    }

    /// Returns the queue's associated CPU index
    #[inline]
    pub fn index(&self) -> usize {
        self.index
    }

    /// "Enters" the scheduler by selecting a first task to execute.
    ///
    /// # Safety
    ///
    /// Only meant to be called once per each CPU when everything's ready.
    pub unsafe fn enter(&self) -> ! {
        let _guard = IrqGuard::acquire();

        (*self.idle.as_ptr()).enter()
    }

    fn pop(&self) -> (Option<Arc<Thread>>, Option<ThreadId>) {
        while let Some(id) = self.queue.pop() {
            let Some(thread) = Thread::get(id) else {
                continue;
            };

            let mut sched = thread.sched.lock();

            assert!(sched.in_queue);
            assert!(core::ptr::eq(self, sched.queue.unwrap()));

            if thread.kill.is_signalled() {
                sched.state = ThreadState::Terminated;
                sched.in_queue = false;
                sched.queue = None;
                thread.set_terminated();
                continue;
            }

            match sched.state {
                ThreadState::Ready => {
                    sched.state = ThreadState::Running;
                    drop(sched);
                    return (Some(thread), Some(id));
                }
                ThreadState::Running => {
                    panic!("Unexpected state: Running ({:?})", id);
                }
                ThreadState::Terminated => {
                    sched.in_queue = false;
                    sched.queue = None;
                    thread.set_terminated();
                }
                ThreadState::Suspended => {
                    sched.queue = None;
                    sched.in_queue = false;
                }
            }
        }

        (None, None)
    }

    fn update_stats(&self, delta: Duration, current: Option<&Arc<Thread>>) {
        let counter = if let Some(current) = current {
            if current.is_user() {
                &self.stats.user
            } else {
                &self.stats.kernel
            }
        } else {
            &self.stats.idle
        };

        let t = delta.as_millis() as u64;

        self.stats.total.fetch_add(t, Ordering::Relaxed);
        counter.fetch_add(t, Ordering::Relaxed);
    }
}

impl Scheduler for CpuQueue {
    type ThreadId = ThreadId;

    fn for_cpu(index: usize) -> &'static Self {
        &QUEUES.get()[index]
    }

    fn for_affinity_mask(mask: u64) -> &'static Self {
        debug_assert_ne!(mask, 0);

        QUEUES
            .get()
            .iter()
            .filter(|c| mask & (1 << c.index) != 0)
            .min_by_key(|c| c.queue.len())
            .unwrap()
    }

    fn local() -> &'static Self {
        CpuImpl::local().scheduler()
    }

    fn is_local(&self) -> bool {
        assert!(ArchitectureImpl::interrupt_mask());
        core::ptr::eq(Self::local(), self)
    }

    fn push(&self, task: Self::ThreadId) {
        self.queue.push(task);
    }

    unsafe fn yield_cpu(&self) -> bool {
        assert!(ArchitectureImpl::interrupt_mask());

        let mut cpu = CpuImpl::<Self>::local();

        let current_id = cpu.current_thread_id();
        let current = current_id.and_then(Thread::get);

        let t = monotonic_time();
        if let Some(t0) = self.last_stats_measure.get() {
            if let Some(dt) = t.checked_sub_time(&t0) {
                self.update_stats(dt, current.as_ref());
            }
        }
        self.last_stats_measure.set(Some(t));

        let drop_current = if let Some(current) = current.as_ref() {
            let mut sched = current.sched.lock();

            let q = sched.queue.unwrap();
            assert!(core::ptr::eq(q, self));
            assert!(sched.in_queue);

            match sched.state {
                ThreadState::Ready => {
                    self.queue.push(current.id);
                    false
                }
                ThreadState::Running => {
                    sched.state = ThreadState::Ready;

                    self.queue.push(current.id);
                    false
                }
                ThreadState::Terminated => {
                    sched.in_queue = false;
                    sched.queue = None;

                    current.set_terminated();
                    true
                }
                ThreadState::Suspended => {
                    sched.queue = None;
                    sched.in_queue = false;
                    false
                }
            }
        } else {
            false
        };

        let (next, next_id) = self.pop();

        let current_ctx = if let Some(current) = current.as_ref() {
            &current.context
        } else {
            &self.idle
        };

        let next_ctx = if let Some(next) = next.as_ref() {
            &next.context
        } else {
            &self.idle
        };

        cpu.set_current_thread_id(next_id);

        let current_ctx = current_ctx.as_ptr();
        let next_ctx = next_ctx.as_ptr();

        if !core::ptr::eq(current_ctx, next_ctx) {
            // Perform the switch
            if drop_current {
                (*next_ctx).switch_and_drop(Arc::into_raw(current.unwrap()) as _);
            } else {
                (*next_ctx).switch(&*current_ctx);
            }

            true
        } else {
            false
        }
    }
}

pub fn stats(index: usize) -> &'static SchedulerStats {
    &QUEUES.get()[index].stats
}

/// Initializes the global queue list
pub fn init_queues(queues: Vec<CpuQueue>) {
    QUEUES.init(queues);
}
