use types::{
    bounded::Priority,
    id::{CoreId, ThreadId},
};

use super::{Scheduler, SchedulerDecision, per_core::PerCoreScheduler};

const MAX_CORES: usize = 8;
const MAX_SCHEDULED_THREADS: usize = 256;
const BALANCE_THRESHOLD: usize = 2;

struct ThreadPlacement {
    core: CoreId,
    priority: Priority,
}

pub struct MultiCoreScheduler {
    cores: [PerCoreScheduler; MAX_CORES],
    num_cores: u8,
    placements: [Option<ThreadPlacement>; MAX_SCHEDULED_THREADS],
}

impl MultiCoreScheduler {
    pub fn new(num_cores: u8) -> Self {
        let cores = core::array::from_fn(|i| PerCoreScheduler::new(CoreId::from_raw(i as u8)));
        Self {
            cores,
            num_cores,
            placements: [const { None }; MAX_SCHEDULED_THREADS],
        }
    }

    pub fn current(&self, core: CoreId) -> Option<ThreadId> {
        self.cores[core.raw() as usize].current()
    }

    pub fn set_current(&mut self, core: CoreId, thread: ThreadId) {
        self.cores[core.raw() as usize].set_current(thread);
    }

    fn thread_count_on(&self, core: CoreId) -> usize {
        self.placements
            .iter()
            .filter(|p| p.as_ref().is_some_and(|tp| tp.core == core))
            .count()
    }

    fn least_loaded_core(&self) -> Option<CoreId> {
        (0..self.num_cores)
            .map(|i| CoreId::from_raw(i))
            .min_by_key(|&c| self.thread_count_on(c))
    }

    fn most_loaded_core(&self) -> Option<CoreId> {
        (0..self.num_cores)
            .map(|i| CoreId::from_raw(i))
            .max_by_key(|&c| self.thread_count_on(c))
    }
}

impl Scheduler for MultiCoreScheduler {
    fn add(&mut self, thread: ThreadId, priority: Priority, core: CoreId) {
        let idx = thread.raw() as usize;
        self.placements[idx] = Some(ThreadPlacement { core, priority });
        self.cores[core.raw() as usize].add(thread, priority);
    }

    fn remove(&mut self, thread: ThreadId) {
        let idx = thread.raw() as usize;
        if let Some(placement) = self.placements[idx].take() {
            let core_idx = placement.core.raw() as usize;
            self.cores[core_idx].remove(thread, placement.priority);
            if self.cores[core_idx].current() == Some(thread) {
                self.cores[core_idx].clear_current();
            }
        }
    }

    fn next(&mut self, core: CoreId) -> Option<ThreadId> {
        self.cores[core.raw() as usize].pick_next()
    }

    fn tick(&mut self, core: CoreId, current: ThreadId) -> SchedulerDecision {
        let idx = current.raw() as usize;
        if let Some(placement) = &self.placements[idx] {
            let current_priority = placement.priority;
            let core_idx = core.raw() as usize;
            if let Some(highest) = self.cores[core_idx].highest_waiting_priority() {
                if highest > current_priority {
                    if let Some(next) = self.cores[core_idx].pick_next() {
                        self.cores[core_idx].add(current, current_priority);
                        return SchedulerDecision::Preempt(next);
                    }
                }
            }
        }
        SchedulerDecision::Continue
    }

    fn balance(&mut self) {
        let (busiest, lightest) = match (self.most_loaded_core(), self.least_loaded_core()) {
            (Some(b), Some(l)) => (b, l),
            _ => return,
        };

        if self.thread_count_on(busiest) > self.thread_count_on(lightest) + BALANCE_THRESHOLD {
            let busiest_idx = busiest.raw() as usize;
            if let Some((thread, priority)) = self.cores[busiest_idx].steal_lowest() {
                let lightest_idx = lightest.raw() as usize;
                self.cores[lightest_idx].add(thread, priority);
                let tid = thread.raw() as usize;
                self.placements[tid] = Some(ThreadPlacement {
                    core: lightest,
                    priority,
                });
            }
        }
    }
}