// ---------------------------------------------------------------------------
// sched — preemptive round-robin kernel tasks.
//
// 16 static task slots.  The 100 Hz PIT tick (vector 32) and the
// software yield vector (0x60) go through `isr_common`, which saves the
// running task's full IntFrame on its own stack and asks the scheduler
// for the next frame pointer — the asm then moves rsp there and iretq's,
// which *is* the context switch.
//
//   - task 0: the shell (boot stack, never allocated by us)
//   - task 1: idle (hlt loop, always ready — the lowest-priority
//     fallback)
//   - tasks 2..: spawn'd, 8 KiB stack + one unmapped guard page below
//
// Preemption is disabled while the scheduler lock (PREEMPT_DEPTH) is
// held: the tick/yield handlers bail out entirely when the depth is
// > 0, so table mutations and klog/serial line writes are atomic with
// respect to switching.  The lock is a counter, not a mutex — re-entrant
// by design (a panic mid-emit increments again; the panic path clears
// it, because a dead task's lock is meaningless).
//
// Task page faults (vector 14) run on a dedicated IST stack (see gdt.rs)
// so a stack overflow inside a guard page kills the task gracefully
// instead of double-faulting the machine.
// ---------------------------------------------------------------------------
use core::arch::asm;
use core::fmt;
use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};

extern crate alloc;
use alloc::string::String;
use alloc::vec::Vec;

use crate::interrupts::IntFrame;
use crate::memory;
use crate::paging;

pub const MAX_TASKS: usize = 16;
pub const SHELL_TID: usize = 0;
pub const IDLE_TID: usize = 1;
const STACK_PAGES: usize = 3; // 1 guard page + 8 KiB stack
const NAME_MAX: usize = 24;
const PANIC_MAX: usize = 64;

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum TaskState {
    Ready,
    Sleeping,
    Blocked,
    Done,
}

struct Task {
    name: [u8; NAME_MAX],
    name_len: u8,
    state: TaskState,
    /// Frame pointer to restore on switch-in (set at spawn; stored on
    /// every switch-away).
    saved_rsp: usize,
    /// Physical base of the stack allocation (0 = none, e.g. the shell).
    stack: usize,
    /// Guard page address (unmapped; 0 = none).
    guard: usize,
    wake_at: u64,
    ticks: u64,
    panic: [u8; PANIC_MAX],
    panic_len: u8,
}

const TASK_NONE: Task = Task {
    name: [0; NAME_MAX],
    name_len: 0,
    state: TaskState::Done,
    saved_rsp: 0,
    stack: 0,
    guard: 0,
    wake_at: 0,
    ticks: 0,
    panic: [0; PANIC_MAX],
    panic_len: 0,
};

static mut TASKS: [Task; MAX_TASKS] = [TASK_NONE; MAX_TASKS];
static mut CURRENT: usize = SHELL_TID;
static mut KERNEL_CS: u16 = 0;
static mut KERNEL_SS: u16 = 0;

static SCHED_READY: AtomicBool = AtomicBool::new(false);
static IN_ISR: AtomicBool = AtomicBool::new(false);
static PREEMPT_DEPTH: AtomicU64 = AtomicU64::new(0);

// ---------------------------------------------------------------------------
// Preemption lock
// ---------------------------------------------------------------------------

/// Disable preemption (re-entrant counter).  The tick/yield handlers
/// skip all scheduling work while the depth is > 0.
pub fn lock() {
    PREEMPT_DEPTH.fetch_add(1, Ordering::Relaxed);
}

/// Re-enable preemption (must pair with `lock`).
pub fn unlock() {
    PREEMPT_DEPTH.fetch_sub(1, Ordering::Relaxed);
}

pub fn in_isr() -> bool {
    IN_ISR.load(Ordering::Relaxed)
}
pub fn set_in_isr(v: bool) {
    IN_ISR.store(v, Ordering::Relaxed);
}
pub fn is_ready() -> bool {
    SCHED_READY.load(Ordering::Relaxed)
}
pub fn current() -> usize {
    unsafe { CURRENT }
}

// ---------------------------------------------------------------------------
// Tick / yield entry points (called from the ISR dispatch)
// ---------------------------------------------------------------------------

/// Timer tick: wake sleepers, account ticks, pick the next task.
/// Receives the interrupted task's frame pointer; returns the frame
/// pointer the asm should switch to (the same one when no switch).
pub fn on_tick(cur: usize) -> usize {
    if !is_ready() || PREEMPT_DEPTH.load(Ordering::Relaxed) > 0 {
        return cur;
    }
    unsafe {
        TASKS[CURRENT].saved_rsp = cur;
        TASKS[CURRENT].ticks += 1;
        let now = crate::interrupts::uptime_ms();
        for t in 0..MAX_TASKS {
            if TASKS[t].state == TaskState::Sleeping && TASKS[t].wake_at <= now {
                TASKS[t].state = TaskState::Ready;
            }
        }
        pick_next(cur)
    }
}

/// Software yield (vector 0x60).
pub fn on_yield(cur: usize) -> usize {
    if !is_ready() || PREEMPT_DEPTH.load(Ordering::Relaxed) > 0 {
        return cur;
    }
    unsafe {
        TASKS[CURRENT].saved_rsp = cur;
        pick_next(cur)
    }
}

/// Round-robin over Ready tasks, skipping the idle task unless nothing
/// else can run.  Returns the next frame pointer (or `cur` to stay).
unsafe fn pick_next(cur: usize) -> usize {
    unsafe {
        for i in 1..MAX_TASKS {
            let idx = (CURRENT + i) % MAX_TASKS;
            if idx == IDLE_TID {
                continue;
            }
            if TASKS[idx].state == TaskState::Ready {
                CURRENT = idx;
                return TASKS[idx].saved_rsp;
            }
        }
        if CURRENT != IDLE_TID && TASKS[IDLE_TID].state == TaskState::Ready {
            CURRENT = IDLE_TID;
            return TASKS[IDLE_TID].saved_rsp;
        }
        cur
    }
}

// ---------------------------------------------------------------------------
// Spawn / lifecycle
// ---------------------------------------------------------------------------

/// Spawn a kernel task running `entry` on its own 8 KiB stack with a
/// guard page below it.  Returns the task id, or None when the table is
/// full (Done slots are reaped first).  `entry` runs inside a
/// trampoline, so returning from it exits the task cleanly.
pub fn spawn(entry: fn(), name: &str) -> Option<usize> {
    lock();
    let id = unsafe {
        // Reap a Done task with a stack, or take a never-used slot.
        let mut slot = None;
        for i in 2..MAX_TASKS {
            if TASKS[i].state == TaskState::Done && TASKS[i].stack != 0 {
                reap(i);
                slot = Some(i);
                break;
            }
            if TASKS[i].state == TaskState::Done && TASKS[i].stack == 0 && TASKS[i].name_len == 0 {
                slot = Some(i);
                break;
            }
        }
        slot?
    };
    let ok = spawn_at(id, entry, name);
    unlock();
    if ok {
        Some(id)
    } else {
        None
    }
}

unsafe fn reap(id: usize) {
    unsafe {
        let t = &mut TASKS[id];
        if t.stack != 0 {
            // Restore the guard page's mapping so the pages are reusable.
            paging::map_page(t.stack, paging::PAGE_WRITABLE);
            memory::free_pages(t.stack, STACK_PAGES);
        }
        *t = TASK_NONE;
    }
}

fn spawn_at(id: usize, entry: fn(), name: &str) -> bool {
    let Some(base) = memory::alloc_pages(STACK_PAGES) else {
        return false;
    };
    // The first page is the guard — unmap it so any stack overflow
    // faults there (caught on the vector-14 IST stack).
    if !paging::unmap_page(base) {
        memory::free_pages(base, STACK_PAGES);
        return false;
    }
    let top = base + STACK_PAGES * memory::PAGE_SIZE;
    let frame = (top - core::mem::size_of::<IntFrame>()) as *mut IntFrame;
    let mut name_buf = [0u8; NAME_MAX];
    let n = name.len().min(NAME_MAX);
    name_buf[..n].copy_from_slice(&name.as_bytes()[..n]);
    unsafe {
        core::ptr::write(
            frame,
            IntFrame {
                r15: 0,
                r14: 0,
                r13: 0,
                r12: 0,
                r11: 0,
                r10: 0,
                r9: 0,
                r8: 0,
                rdi: 0,
                rsi: 0,
                rbp: 0,
                rbx: 0,
                rdx: 0,
                rcx: entry as usize as u64, // MS x64 ABI: first arg
                rax: 0,
                vector: 0,
                error: 0,
                rip: trampoline as *const () as u64,
                cs: KERNEL_CS as u64,
                rflags: 0x202, // IF | reserved bit 1
                // QEMU's iretq always pops RSP/SS from the frame tail.
                // rsp = the stack top so the first iretq lands with a
                // valid stack pointer (the trampoline then descends from
                // it normally); ss mirrors the kernel data selector.
                rsp: top as u64,
                ss: KERNEL_SS as u64,
            },
        );
        TASKS[id] = Task {
            name: name_buf,
            name_len: n as u8,
            state: TaskState::Ready,
            saved_rsp: frame as usize,
            stack: base,
            guard: base,
            wake_at: 0,
            ticks: 0,
            panic: [0; PANIC_MAX],
            panic_len: 0,
        };
    }
    true
}

/// Entry trampoline: run the task body, then exit the task when it
/// returns.  Runs with the synthetic frame's registers (rcx = entry).
unsafe extern "C" fn trampoline(entry: usize) -> ! {
    let f: fn() = unsafe { core::mem::transmute(entry as *const ()) };
    f();
    task_exit()
}

fn task_exit() -> ! {
    lock();
    unsafe {
        TASKS[CURRENT].state = TaskState::Done;
    }
    unlock();
    yield_int();
    unreachable!()
}

/// Block the current task for `ms` milliseconds.  The shell uses this
/// for `sleep`; background tasks keep running meanwhile.
pub fn sleep_ms(ms: u64) {
    lock();
    unsafe {
        TASKS[CURRENT].state = TaskState::Sleeping;
        TASKS[CURRENT].wake_at = crate::interrupts::uptime_ms() + ms;
    }
    unlock();
    yield_int();
}

/// Cooperatively yield the rest of the current quantum.
pub fn yield_now() {
    yield_int();
}

/// Software-interrupt the scheduler (vector 0x60).  Net stack usage is
/// zero (the frame is popped by iretq), so no `nostack` — the compiler
/// must not use the red zone around the interrupt.
#[inline(never)]
fn yield_int() {
    unsafe {
        asm!("int 0x60");
    }
}

// ---------------------------------------------------------------------------
// Task panics and faults
// ---------------------------------------------------------------------------

/// Kill the current task with a panic message: logs it, marks the task
/// Done, drops its scheduler lock (a dead task's lock is meaningless),
/// and switches away forever.  Called by the panic hook (task context)
/// and by the page-fault handler (IST stack, ISR context).
pub fn abort_current(msg: &str) -> ! {
    unsafe {
        let id = CURRENT;
        let n = msg.len().min(PANIC_MAX);
        TASKS[id].panic[..n].copy_from_slice(&msg.as_bytes()[..n]);
        TASKS[id].panic_len = n as u8;
        TASKS[id].state = TaskState::Done;
        crate::serial_println!(
            "sched: task {} \"{}\" panicked: {} (rsp_saved={:#x} stack={:#x} guard={:#x})",
            id,
            name_str(id),
            msg,
            TASKS[id].saved_rsp,
            TASKS[id].stack,
            TASKS[id].guard
        );
        PREEMPT_DEPTH.store(0, Ordering::Relaxed);
    }
    yield_int();
    unreachable!()
}

/// Vector 14 (page fault) handler — runs on the dedicated IST stack.
/// A fault in a real task (id >= 2) kills the task gracefully; anything
/// else (shell, idle, pre-scheduler) is a kernel bug: halt.
pub fn page_fault(cr2: u64, frame_ptr: usize) -> usize {
    if !is_ready() || unsafe { CURRENT } <= IDLE_TID {
        crate::serial_println!("pgfault: kernel page fault at {:#x} — halting", cr2);
        crate::interrupts::halt();
    }
    let g = unsafe { TASKS[CURRENT].guard };
    let in_guard = g != 0 && cr2 >= g as u64 && cr2 < (g + memory::PAGE_SIZE) as u64;
    let rip = unsafe { (frame_ptr as *const IntFrame).read().rip };
    let mut buf = [0u8; 96];
    let n = {
        let mut f = FmtBuf { b: &mut buf, n: 0 };
        let _ = fmt::write(
            &mut f,
            core::format_args!(
                "{} at {:#x} (rip={:#x})",
                if in_guard {
                    "stack overflow (guard page)"
                } else {
                    "page fault"
                },
                cr2,
                rip
            ),
        );
        f.n
    };
    abort_current(core::str::from_utf8(&buf[..n]).unwrap_or("task fault"))
}

// ---------------------------------------------------------------------------
// Introspection (shell: jobs / kill)
// ---------------------------------------------------------------------------

/// Snapshot the task table for `jobs` — one line per used slot.
pub fn jobs() -> Vec<String> {
    let mut v = Vec::new();
    lock();
    unsafe {
        for i in 0..MAX_TASKS {
            let t = &TASKS[i];
            if t.name_len == 0 {
                continue; // never used
            }
            let state = match t.state {
                TaskState::Ready => "ready",
                TaskState::Sleeping => "sleeping",
                TaskState::Blocked => "blocked",
                TaskState::Done => {
                    if t.panic_len > 0 {
                        "done (panicked)"
                    } else {
                        "done"
                    }
                }
            };
            let mut line = alloc::format!(
                "  {:>2}  {:<12} {:<15} {:>6} ticks",
                i,
                name_str(i),
                state,
                t.ticks
            );
            if t.panic_len > 0 {
                let m = core::str::from_utf8(&t.panic[..t.panic_len as usize]).unwrap_or("?");
                line.push_str("  ");
                line.push_str(m);
            }
            v.push(line);
        }
    }
    unlock();
    v
}

/// Mark a task Done so it stops running and gets reaped at the next
/// spawn.  The shell (0) and idle (1) are protected.
pub fn kill(id: usize) -> Result<(), &'static str> {
    if id == SHELL_TID {
        return Err("cannot kill the shell (task 0)");
    }
    if id == IDLE_TID {
        return Err("cannot kill the idle task (task 1)");
    }
    if id >= MAX_TASKS {
        return Err("no such task");
    }
    lock();
    let alive = unsafe {
        let t = &TASKS[id];
        t.name_len != 0 && t.state != TaskState::Done
    };
    if !alive {
        unlock();
        return Err("no such task");
    }
    unsafe {
        TASKS[id].state = TaskState::Done;
    }
    unlock();
    Ok(())
}

pub fn name_str(id: usize) -> &'static str {
    unsafe {
        let t = &TASKS[id];
        if t.name_len == 0 {
            return "";
        }
        core::str::from_utf8(&t.name[..t.name_len as usize]).unwrap_or("")
    }
}

// ---------------------------------------------------------------------------
// Init
// ---------------------------------------------------------------------------

/// Set up the task table: task 0 is the shell (boot stack), task 1 is
/// the idle `hlt` fallback.  Called after `interrupts::init` (the tick
/// must be live so `sleep`/wake scheduling works).
pub fn init() {
    unsafe {
        let cs64: u64;
        asm!("mov {0}, cs", out(reg) cs64, options(nostack, preserves_flags));
        KERNEL_CS = cs64 as u16;
        let ss64: u64;
        asm!("mov {0}, ss", out(reg) ss64, options(nostack, preserves_flags));
        KERNEL_SS = ss64 as u16;

        TASKS[SHELL_TID] = Task {
            name: {
                let mut n = [0u8; 24];
                n[..5].copy_from_slice(b"shell");
                n
            },
            name_len: 5,
            state: TaskState::Ready,
            saved_rsp: 0, // captured on the first switch-away
            stack: 0,
            guard: 0,
            wake_at: 0,
            ticks: 0,
            panic: [0; PANIC_MAX],
            panic_len: 0,
        };
        spawn_at(IDLE_TID, idle_main, "idle");
        SCHED_READY.store(true, Ordering::Relaxed);
    }
    crate::serial_println!(
        "sched: ready — {} slots, shell=task 0, idle=task 1",
        MAX_TASKS
    );
    crate::serial_println!(
        "sched: trampoline={:#x} idle={:#x} emit={:#x}",
        trampoline as *const () as usize,
        idle_main as *const () as usize,
        crate::klog::emit as *const () as usize
    );
}

fn idle_main() {
    loop {
        unsafe {
            asm!("hlt", options(nostack, preserves_flags));
        }
    }
}

struct FmtBuf<'a> {
    b: &'a mut [u8],
    n: usize,
}

impl fmt::Write for FmtBuf<'_> {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        let m = s.len().min(self.b.len() - self.n);
        self.b[self.n..self.n + m].copy_from_slice(&s.as_bytes()[..m]);
        self.n += m;
        Ok(())
    }
}