use crate::cap::cnode;
use crate::cap::object::ObjectTag;
use crate::cap::pool::POOL;
use crate::cap::table::Rights;
use crate::error::KernelError;
use crate::proc::PROCESSES;
use crate::proc::context::CpuContext;
use crate::syscall::{SyscallResult, try_syscall};
use crate::types::{Pid, Priority};
use lancer_core::header::NONE_SENTINEL;
use lancer_core::object_layout::SchedContextObject;

pub fn sys_sched_attach(ctx: &mut CpuContext) {
    let sched_addr = ctx.rdi;
    let target_pid_raw = try_syscall!(ctx, super::u32_from_reg(ctx.rsi));
    let pid = crate::arch::syscall::current_pid();

    let target_pid = match Pid::try_new(target_pid_raw) {
        Some(p) => p,
        None => {
            ctx.rax = SyscallResult::error(KernelError::InvalidParameter).raw();
            return;
        }
    };

    let mut ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        try_syscall!(
            ctx,
            cnode::resolve_caller_validate(
                pid,
                sched_addr,
                ObjectTag::SchedContext,
                Rights::WRITE,
                &ptable,
                &pool
            )
        )
    };

    if target_pid != pid {
        let pool = POOL.lock_after(&ptable);
        let has_proc_cap = {
            let (cnode_id, cnode_gen, _, _, _) =
                try_syscall!(ctx, cnode::cnode_coords(pid, &ptable));
            let mut found = false;
            let _ = cnode::walk_cnode_slots(&pool, cnode_id, cnode_gen, |slot| {
                if let crate::cap::table::CapSlot::Active(cap) = slot
                    && cap.tag() == ObjectTag::Process
                    && cap.rights().contains(Rights::WRITE)
                    && pool
                        .read_as::<lancer_core::object_layout::ProcessObject>(
                            cap.phys(),
                            cap.generation(),
                        )
                        .ok()
                        .and_then(|p| Pid::try_new(p.pid))
                        .is_some_and(|p| p == target_pid)
                {
                    found = true;
                }
            });
            found
        };
        if !has_proc_cap {
            ctx.rax = SyscallResult::error(KernelError::PermissionDenied).raw();
            return;
        }
    }

    let priority = {
        let mut pool = POOL.lock_after(&ptable);
        match pool.write_as::<SchedContextObject>(cap.phys(), cap.generation()) {
            Ok(sc) => {
                let old_attached = Pid::try_new(sc.attached_pid);
                match old_attached {
                    Some(old_pid) if old_pid != target_pid => {
                        if old_pid != pid {
                            ctx.rax = SyscallResult::error(KernelError::PermissionDenied).raw();
                            return;
                        }
                        if let Some(old_proc) = ptable.get_mut(old_pid) {
                            old_proc.detach_sched_context();
                        }
                    }
                    _ => {}
                }
                sc.attached_pid = target_pid.raw();
                let now_us = crate::wcet::tsc::cycles_to_ns(crate::wcet::tsc::read_tsc()) / 1000;
                sc.remaining_us = sc.budget_us;
                sc.replenish_at = now_us.saturating_add(sc.period_us);
                Priority::new(sc.priority)
            }
            Err(e) => {
                ctx.rax = SyscallResult::error(e).raw();
                return;
            }
        }
    };

    ctx.rax = match ptable.get_mut(target_pid) {
        Some(target) => {
            let old_sc = target.sched_context();
            target.attach_sched_context(cap.phys(), cap.generation(), priority);

            if let Some((old_id, old_gen)) = old_sc {
                let mut pool = POOL.lock_after(&ptable);
                if let Ok(sc) = pool.write_as::<SchedContextObject>(old_id, old_gen) {
                    sc.attached_pid = NONE_SENTINEL;
                }
            }
            SyscallResult::ok().raw()
        }
        None => SyscallResult::error(KernelError::InvalidObject).raw(),
    };
}

pub fn sys_sched_configure(ctx: &mut CpuContext) {
    let sched_addr = ctx.rdi;
    let budget_us = ctx.rsi;
    let period_us = ctx.rdx;
    let priority_raw = try_syscall!(ctx, super::u8_from_reg(ctx.r10));
    let pid = crate::arch::syscall::current_pid();

    const MIN_PERIOD_US: u64 = 1_000;
    const MAX_PERIOD_US: u64 = 3_600_000_000;

    let priority = match Priority::new_user(priority_raw) {
        Some(p) => p,
        None => {
            ctx.rax = SyscallResult::error(KernelError::InvalidParameter).raw();
            return;
        }
    };

    if budget_us == 0
        || period_us < MIN_PERIOD_US
        || budget_us > period_us
        || period_us > MAX_PERIOD_US
    {
        ctx.rax = SyscallResult::error(KernelError::InvalidParameter).raw();
        return;
    }

    let ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        try_syscall!(
            ctx,
            cnode::resolve_caller_validate(
                pid,
                sched_addr,
                ObjectTag::SchedContext,
                Rights::WRITE,
                &ptable,
                &pool,
            )
        )
    };

    let mut pool = POOL.lock_after(&ptable);
    ctx.rax = match pool.write_as::<SchedContextObject>(cap.phys(), cap.generation()) {
        Ok(sc) => {
            sc.budget_us = budget_us;
            sc.period_us = period_us;
            sc.remaining_us = budget_us;
            sc.priority = priority.raw();
            SyscallResult::ok().raw()
        }
        Err(e) => SyscallResult::error(e).raw(),
    };
}

pub fn sys_sched_yield(ctx: &mut CpuContext) {
    ctx.rax = SyscallResult::ok().raw();
    crate::sched::schedule(ctx);
}