use types::{
    bounded::{CapSlot, Priority},
    id::{CoreId, ProcessId, ThreadId},
    time::TimeSlice,
};

use crate::error::KernelError;

#[cfg(test)]
#[path = "tests.rs"]
mod tests;

pub struct Thread {
    id: ThreadId,
    process: ProcessId,
    state: ThreadState,
    saved_context: SavedContext,
    priority: Priority,
    time_slice: TimeSlice,
    core_affinity: Option<CoreId>,
}

pub enum ThreadState {
    Ready,
    Running(CoreId),
    Blocked(BlockReason),
    Dead,
}

pub struct SavedContext {
    registers: [usize; 32],
}

pub const PC_REGISTER: usize = 0;
pub const SP_REGISTER: usize = 1;

impl SavedContext {
    pub fn empty() -> Self {
        Self { registers: [0; 32] }
    }

    pub fn with_entry_and_stack(
        entry: types::address::EntryPoint,
        stack: types::address::StackPointer,
    ) -> Self {
        let mut ctx = Self::empty();
        ctx.registers[PC_REGISTER] = entry.raw();
        ctx.registers[SP_REGISTER] = stack.raw();
        ctx
    }

    pub fn program_counter(&self) -> usize {
        self.registers[PC_REGISTER]
    }

    pub fn stack_pointer(&self) -> usize {
        self.registers[SP_REGISTER]
    }
}

pub enum BlockReason {
    Sending(CapSlot),
    Receiving(CapSlot),
    WaitingForInterrupt(CapSlot),
    WaitingAny,
}

impl Thread {
    pub fn new(
        id: ThreadId,
        process: ProcessId,
        priority: Priority,
        time_slice: TimeSlice,
        core_affinity: Option<CoreId>,
    ) -> Self {
        Self {
            id,
            process,
            state: ThreadState::Ready,
            saved_context: SavedContext::empty(),
            priority,
            time_slice,
            core_affinity,
        }
    }

    pub fn with_context(
        id: ThreadId,
        process: ProcessId,
        priority: Priority,
        time_slice: TimeSlice,
        core_affinity: Option<CoreId>,
        context: SavedContext,
    ) -> Self {
        Self {
            id,
            process,
            state: ThreadState::Ready,
            saved_context: context,
            priority,
            time_slice,
            core_affinity,
        }
    }

    pub fn saved_context(&self) -> &SavedContext {
        &self.saved_context
    }

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

    pub fn process(&self) -> ProcessId {
        self.process
    }

    pub fn state(&self) -> &ThreadState {
        &self.state
    }

    pub fn priority(&self) -> Priority {
        self.priority
    }

    pub fn time_slice(&self) -> &TimeSlice {
        &self.time_slice
    }

    pub fn time_slice_mut(&mut self) -> &mut TimeSlice {
        &mut self.time_slice
    }

    pub fn core_affinity(&self) -> Option<CoreId> {
        self.core_affinity
    }

    pub fn set_running(&mut self, core: CoreId) -> Result<(), KernelError> {
        match self.state {
            ThreadState::Ready => {
                self.state = ThreadState::Running(core);
                Ok(())
            }
            _ => Err(KernelError::InvalidStateTransition(self.id)),
        }
    }

    pub fn block(&mut self, reason: BlockReason) -> Result<(), KernelError> {
        match self.state {
            ThreadState::Running(_) => {
                self.state = ThreadState::Blocked(reason);
                Ok(())
            }
            _ => Err(KernelError::InvalidStateTransition(self.id)),
        }
    }

    pub fn unblock(&mut self) -> Result<(), KernelError> {
        match self.state {
            ThreadState::Blocked(_) => {
                self.state = ThreadState::Ready;
                Ok(())
            }
            _ => Err(KernelError::InvalidStateTransition(self.id)),
        }
    }

    pub fn preempt(&mut self) -> Result<(), KernelError> {
        match self.state {
            ThreadState::Running(_) => {
                self.state = ThreadState::Ready;
                Ok(())
            }
            _ => Err(KernelError::InvalidStateTransition(self.id)),
        }
    }

    pub fn exit(&mut self) -> Result<(), KernelError> {
        match self.state {
            ThreadState::Running(_) => {
                self.state = ThreadState::Dead;
                Ok(())
            }
            _ => Err(KernelError::InvalidStateTransition(self.id)),
        }
    }
}

pub const MAX_THREADS: usize = 256;

pub struct ThreadTable {
    threads: [Option<Thread>; MAX_THREADS],
}

impl ThreadTable {
    pub fn new() -> Self {
        Self {
            threads: [const { None }; MAX_THREADS],
        }
    }

    pub fn spawn(
        &mut self,
        process: ProcessId,
        priority: Priority,
        time_slice: TimeSlice,
        core_affinity: Option<CoreId>,
    ) -> Result<ThreadId, KernelError> {
        let idx = self
            .threads
            .iter()
            .position(|s| s.is_none())
            .ok_or(KernelError::TooManyThreads)?;

        let id = ThreadId::from_raw(idx as u16);
        self.threads[idx] = Some(Thread::new(id, process, priority, time_slice, core_affinity));
        Ok(id)
    }

    pub fn spawn_with_context(
        &mut self,
        process: ProcessId,
        priority: Priority,
        time_slice: TimeSlice,
        core_affinity: Option<CoreId>,
        context: SavedContext,
    ) -> Result<ThreadId, KernelError> {
        let idx = self
            .threads
            .iter()
            .position(|s| s.is_none())
            .ok_or(KernelError::TooManyThreads)?;

        let id = ThreadId::from_raw(idx as u16);
        self.threads[idx] = Some(Thread::with_context(id, process, priority, time_slice, core_affinity, context));
        Ok(id)
    }

    pub fn get(&self, id: ThreadId) -> Result<&Thread, KernelError> {
        let idx = id.raw() as usize;
        self.threads
            .get(idx)
            .and_then(|s| s.as_ref())
            .ok_or(KernelError::InvalidThread(id))
    }

    pub fn get_mut(&mut self, id: ThreadId) -> Result<&mut Thread, KernelError> {
        let idx = id.raw() as usize;
        self.threads
            .get_mut(idx)
            .and_then(|s| s.as_mut())
            .ok_or(KernelError::InvalidThread(id))
    }

    pub fn destroy(&mut self, id: ThreadId) -> Result<Thread, KernelError> {
        let idx = id.raw() as usize;
        self.threads
            .get_mut(idx)
            .and_then(|s| s.take())
            .ok_or(KernelError::InvalidThread(id))
    }

    pub fn threads_for_process(&self, process: ProcessId) -> impl Iterator<Item = &Thread> {
        self.threads
            .iter()
            .filter_map(|s| s.as_ref())
            .filter(move |t| t.process() == process)
    }
}