//! Task — структура, описывающая поток выполнения (процесс или поток ядра).

use alloc::alloc::{alloc_zeroed, dealloc, Layout};
use alloc::string::String;
use alloc::sync::Arc;
use core::sync::atomic::{AtomicU64, AtomicU8, Ordering};
use spin::Mutex;

/// Размер стэка ядра для каждого потока (16 KiB = 4 страницы).
pub const KERNEL_STACK_SIZE: usize = 16 * 1024;

/// Уникальный идентификатор задачи.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct TaskId(pub u64);

impl TaskId {
    pub fn new() -> Self {
        static COUNTER: AtomicU64 = AtomicU64::new(1);
        TaskId(COUNTER.fetch_add(1, Ordering::Relaxed))
    }

    pub fn as_u64(self) -> u64 {
        self.0
    }
}

impl core::fmt::Display for TaskId {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "Task#{}", self.0)
    }
}

/// Состояние задачи.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum TaskState {
    /// Готова к выполнению, ожидает в очереди планировщика.
    Ready = 0,
    /// Выполняется на CPU прямо сейчас.
    Running = 1,
    /// Заблокирована (ждёт таймера, IPC, семафора).
    Blocked = 2,
    /// Завершилась, ожидает очистки.
    Dead = 3,
}

impl From<u8> for TaskState {
    fn from(val: u8) -> Self {
        match val {
            0 => TaskState::Ready,
            1 => TaskState::Running,
            2 => TaskState::Blocked,
            _ => TaskState::Dead,
        }
    }
}

/// Приоритет задачи (0 — наивысший, 7 — наименьший).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Priority(pub u8);

impl Priority {
    pub const KERNEL:   Self = Priority(0);
    pub const REALTIME: Self = Priority(1);
    pub const HIGH:     Self = Priority(2);
    pub const NORMAL:   Self = Priority(4);
    pub const LOW:      Self = Priority(6);
    pub const IDLE:     Self = Priority(7);
}

/// Выделенный стэк ядра.
pub struct KernelStack {
    ptr: *mut u8,
    layout: Layout,
}

unsafe impl Send for KernelStack {}
unsafe impl Sync for KernelStack {}

impl KernelStack {
    /// Выделить новый стек ядра размером 16 KiB.
    pub fn new() -> Self {
        let layout = Layout::from_size_align(KERNEL_STACK_SIZE, 4096).unwrap();
        let ptr = unsafe { alloc_zeroed(layout) };
        assert!(!ptr.is_null(), "KernelStack: OOM allocating kernel stack");
        KernelStack { ptr, layout }
    }

    /// Верхушка стека (адрес для начала роста стека вниз).
    pub fn top(&self) -> usize {
        (self.ptr as usize) + KERNEL_STACK_SIZE
    }
}

impl Drop for KernelStack {
    fn drop(&mut self) {
        unsafe {
            dealloc(self.ptr, self.layout);
        }
    }
}

/// Задача (поток выполнения).
pub struct Task {
    pub id:        TaskId,
    pub name:      String,
    pub priority:  Priority,
    pub state:     AtomicU8,
    pub rsp:       Mutex<usize>,
    pub stack:     Option<KernelStack>,
    pub ticks_run: AtomicU64,
}

impl Task {
    /// Создать структуру для текущего главного потока ядра (он уже имеет свой стек).
    pub fn new_main(name: &str) -> Arc<Self> {
        Arc::new(Task {
            id: TaskId(0),
            name: String::from(name),
            priority: Priority::NORMAL,
            state: AtomicU8::new(TaskState::Running as u8),
            rsp: Mutex::new(0),
            stack: None,
            ticks_run: AtomicU64::new(0),
        })
    }

    /// Создать новый поток ядра со своим выделенным стеком.
    pub fn new_kernel_thread(
        name: &str,
        entry_fn: extern "C" fn(usize),
        arg: usize,
        priority: Priority,
    ) -> Arc<Self> {
        let stack = KernelStack::new();
        let stack_top = stack.top();

        // Формируем начальный стек для __switch_context:
        // top - 8:  arg (usize)
        // top - 16: entry_fn (usize)
        // top - 24: thread_trampoline RIP
        // top - 32: rbp (0)
        // top - 40: rbx (0)
        // top - 48: r12 (0)
        // top - 56: r13 (0)
        // top - 64: r14 (0)
        // top - 72: r15 (0)
        let initial_rsp = unsafe {
            let sp = stack_top as *mut usize;
            *sp.offset(-1) = arg;
            *sp.offset(-2) = entry_fn as usize;
            *sp.offset(-3) = crate::arch::x86_64::context::thread_trampoline as *const () as usize;
            *sp.offset(-4) = 0; // rbp
            *sp.offset(-5) = 0; // rbx
            *sp.offset(-6) = 0; // r12
            *sp.offset(-7) = 0; // r13
            *sp.offset(-8) = 0; // r14
            *sp.offset(-9) = 0; // r15
            (sp.offset(-9)) as usize
        };

        Arc::new(Task {
            id: TaskId::new(),
            name: String::from(name),
            priority,
            state: AtomicU8::new(TaskState::Ready as u8),
            rsp: Mutex::new(initial_rsp),
            stack: Some(stack),
            ticks_run: AtomicU64::new(0),
        })
    }

    pub fn get_state(&self) -> TaskState {
        TaskState::from(self.state.load(Ordering::Relaxed))
    }

    pub fn set_state(&self, state: TaskState) {
        self.state.store(state as u8, Ordering::Relaxed);
    }
}