use crate::cap::object::PidQueue;
use crate::cap::pool::POOL;
use crate::cap::table::CapRef;
use crate::ipc::{AlwaysBlocked, CallResult, IpcOutcome, IpcResult, message};
use crate::proc::{BlockedReason, ProcessManager};
use crate::types::{BlockedPid, MAX_PIDS, Pid};
use lancer_core::header::NONE_SENTINEL;
use lancer_core::object_layout::EndpointObject;

pub(crate) fn load_senders(ep: &EndpointObject) -> PidQueue {
    PidQueue::from_repr_c(ep.sender_head, ep.sender_tail, ep.sender_len)
}

pub(crate) fn load_receivers(ep: &EndpointObject) -> PidQueue {
    PidQueue::from_repr_c(ep.receiver_head, ep.receiver_tail, ep.receiver_len)
}

pub(crate) fn store_senders(ep: &mut EndpointObject, q: &PidQueue) {
    let (h, t, l) = q.to_repr_c();
    ep.sender_head = h;
    ep.sender_tail = t;
    ep.sender_len = l;
}

pub(crate) fn store_receivers(ep: &mut EndpointObject, q: &PidQueue) {
    let (h, t, l) = q.to_repr_c();
    ep.receiver_head = h;
    ep.receiver_tail = t;
    ep.receiver_len = l;
}

pub(crate) fn load_holder(ep: &EndpointObject) -> Option<Pid> {
    Pid::try_new(ep.holder)
}

pub(crate) fn store_holder(ep: &mut EndpointObject, holder: Option<Pid>) {
    ep.holder = holder.map_or(NONE_SENTINEL, |p| p.raw());
}

fn assert_queue_valid(queue: &PidQueue, ptable: &ProcessManager) {
    let valid = queue.validate(|pid| ptable[pid].next_ipc);
    debug_assert!(valid, "PidQueue invariant violation detected");
}

pub(crate) fn enqueue(
    queue: &mut PidQueue,
    blocked: BlockedPid,
    ptable: &mut ProcessManager,
) -> Result<(), (crate::error::KernelError, BlockedPid)> {
    let pid = blocked.pid();
    if ptable[pid].next_ipc.is_some() {
        return Err((crate::error::KernelError::BadState, blocked));
    }
    queue.check_capacity().map_err(|e| (e, blocked))?;
    match queue.tail {
        Some(tail) => {
            ptable[tail].next_ipc = Some(pid);
            queue.tail = Some(pid);
        }
        None => {
            queue.head = Some(pid);
            queue.tail = Some(pid);
        }
    }
    let r = queue.inc_len();
    debug_assert!(r.is_ok(), "inc_len cannot fail after check_capacity");
    assert_queue_valid(queue, ptable);
    Ok(())
}

pub(crate) fn dequeue(queue: &mut PidQueue, ptable: &mut ProcessManager) -> Option<BlockedPid> {
    queue.head.map(|pid| {
        queue.head = ptable[pid].next_ipc;
        ptable[pid].next_ipc = None;
        if queue.head.is_none() {
            queue.tail = None;
        }
        queue.dec_len();
        assert_queue_valid(queue, ptable);
        ptable[pid].blocked_proof()
    })
}

pub(crate) fn dequeue_genuine_receiver(
    queue: &mut PidQueue,
    ptable: &mut ProcessManager,
) -> Option<BlockedPid> {
    let candidate = core::iter::successors(queue.head, |&cur| ptable[cur].next_ipc)
        .take(MAX_PIDS)
        .find(|&pid| {
            !matches!(
                ptable[pid].blocked_reason(),
                Some(BlockedReason::Calling(_, _))
            )
        })?;
    remove_pid(queue, candidate, ptable);
    Some(ptable[candidate].blocked_proof())
}

fn remove_pid(queue: &mut PidQueue, target: Pid, ptable: &mut ProcessManager) {
    let Some(h) = queue.head else { return };

    if h == target {
        queue.head = ptable[target].next_ipc;
        ptable[target].next_ipc = None;
        if queue.head.is_none() {
            queue.tail = None;
        }
        queue.dec_len();
        assert_queue_valid(queue, ptable);
        return;
    }

    let Some(prev) = core::iter::successors(queue.head, |&cur| ptable[cur].next_ipc)
        .take(MAX_PIDS)
        .find(|&cur| ptable[cur].next_ipc == Some(target))
    else {
        return;
    };

    let next_after_target = ptable[target].next_ipc;
    ptable[prev].next_ipc = next_after_target;
    ptable[target].next_ipc = None;
    if next_after_target.is_none() {
        queue.tail = Some(prev);
    }
    queue.dec_len();
    assert_queue_valid(queue, ptable);
}

pub fn do_send(cap: &CapRef, sender_pid: Pid, ptable: &mut ProcessManager) -> IpcResult<()> {
    let sender_msg = ptable.exec(sender_pid).unwrap().ipc_message;

    let blocked_recv = {
        let mut pool = POOL.lock();
        let ep = pool.write_as::<EndpointObject>(cap.phys(), cap.generation())?;
        let mut receivers = load_receivers(ep);
        let dequeued = dequeue_genuine_receiver(&mut receivers, ptable);
        if let Some(br) = &dequeued {
            store_holder(ep, Some(br.pid()));
        }
        store_receivers(ep, &receivers);
        dequeued
    };

    match blocked_recv {
        Some(blocked_recv) => {
            let recv_pid = blocked_recv.pid();
            let sender_prio = ptable[sender_pid].effective_priority();

            ptable.unblock_and_enqueue(recv_pid, blocked_recv)?;
            let recv_exec = ptable.exec_mut(recv_pid).unwrap();
            recv_exec.ipc_message = sender_msg;
            recv_exec.ipc_badge = sender_pid.raw() as u64;
            message::inject_into_context(&mut recv_exec.saved_context, &sender_msg);
            recv_exec.saved_context.rax = sender_pid.raw() as u64;
            recv_exec.seal_context();

            crate::sched::boost_effective(ptable, recv_pid, sender_prio);

            Ok(IpcOutcome::Done(()))
        }
        None => {
            let mut pool = POOL.lock();
            let (holder, donor_prio) = {
                let ep = pool.write_as::<EndpointObject>(cap.phys(), cap.generation())?;
                let mut senders = load_senders(ep);
                let blocked_sender = ptable[sender_pid]
                    .block_on(BlockedReason::Sending(cap.phys(), cap.generation()))?;
                enqueue(&mut senders, blocked_sender, ptable).map_err(|(e, proof)| {
                    let _ = ptable.unblock_and_enqueue(sender_pid, proof);
                    e
                })?;
                store_senders(ep, &senders);
                let holder = load_holder(ep);
                (holder, ptable[sender_pid].effective_priority())
            };
            if let Some(h) = holder {
                crate::sched::propagate_priority(ptable, &pool, h, donor_prio);
            }
            Ok(IpcOutcome::Blocked)
        }
    }
}

fn dequeue_sender(
    cap: &CapRef,
    recv_pid: Pid,
    ptable: &mut ProcessManager,
) -> Result<Option<BlockedPid>, crate::error::KernelError> {
    let mut pool = POOL.lock();
    let ep = pool.write_as::<EndpointObject>(cap.phys(), cap.generation())?;
    let mut senders = load_senders(ep);
    let dequeued = dequeue(&mut senders, ptable);
    if dequeued.is_some() {
        store_holder(ep, Some(recv_pid));
    }
    store_senders(ep, &senders);
    Ok(dequeued)
}

fn handle_dequeued_sender(
    cap: &CapRef,
    recv_pid: Pid,
    blocked_sender: BlockedPid,
    ptable: &mut ProcessManager,
) -> IpcResult<Pid> {
    let sender_pid = blocked_sender.pid();
    let sender_msg = ptable.exec(sender_pid).unwrap().ipc_message;
    let sender_prio = ptable[sender_pid].effective_priority();

    match ptable[sender_pid].blocked_reason() {
        Some(BlockedReason::Calling(_, _)) => {
            let mut pool = POOL.lock();
            match pool.write_as::<EndpointObject>(cap.phys(), cap.generation()) {
                Ok(ep) => {
                    let mut receivers = load_receivers(ep);
                    match enqueue(&mut receivers, blocked_sender, ptable) {
                        Ok(()) => {
                            store_receivers(ep, &receivers);
                            let recv_exec = ptable.exec_mut(recv_pid).unwrap();
                            recv_exec.ipc_message = sender_msg;
                            recv_exec.ipc_badge = sender_pid.raw() as u64;
                            recv_exec.reply_target = Some(sender_pid);
                        }
                        Err((e, proof)) => {
                            store_receivers(ep, &receivers);
                            crate::show!(ipc, error, "recv re-enqueue caller failed {:?}", e);
                            let sender_exec = ptable.exec_mut(sender_pid).unwrap();
                            sender_exec.saved_context.rax =
                                crate::error::KernelError::ResourceExhausted.to_errno() as u64;
                            sender_exec.seal_context();
                            {
                                let r = ptable.unblock_and_enqueue(sender_pid, proof);
                                debug_assert!(r.is_ok());
                            }
                        }
                    }
                }
                Err(e) => {
                    let sender_exec = ptable.exec_mut(sender_pid).unwrap();
                    sender_exec.saved_context.rax = e.to_errno() as u64;
                    sender_exec.seal_context();
                    {
                        let r = ptable.unblock_and_enqueue(sender_pid, blocked_sender);
                        debug_assert!(r.is_ok());
                    }
                }
            }
        }
        Some(BlockedReason::Sending(_, _))
        | Some(BlockedReason::Receiving(_, _))
        | Some(BlockedReason::WaitingNotification(_, _))
        | None => {
            let recv_exec = ptable.exec_mut(recv_pid).unwrap();
            recv_exec.ipc_message = sender_msg;
            recv_exec.ipc_badge = sender_pid.raw() as u64;
            ptable.unblock_and_enqueue(sender_pid, blocked_sender)?;
        }
    }

    {
        let pool = POOL.lock();
        match pool.read_as::<EndpointObject>(cap.phys(), cap.generation()) {
            Ok(ep) => {
                let senders = load_senders(ep);
                crate::sched::recalculate_effective(ptable, recv_pid, &senders);
                crate::sched::boost_effective(ptable, recv_pid, sender_prio);
            }
            Err(_) => {
                crate::sched::reset_effective(ptable, recv_pid);
                crate::sched::boost_effective(ptable, recv_pid, sender_prio);
            }
        }
    }

    Ok(IpcOutcome::Done(sender_pid))
}

pub fn do_recv(cap: &CapRef, recv_pid: Pid, ptable: &mut ProcessManager) -> IpcResult<Pid> {
    match dequeue_sender(cap, recv_pid, ptable)? {
        Some(blocked_sender) => handle_dequeued_sender(cap, recv_pid, blocked_sender, ptable),
        None => {
            let mut pool = POOL.lock();
            let ep = pool.write_as::<EndpointObject>(cap.phys(), cap.generation())?;
            let mut receivers = load_receivers(ep);
            let blocked_recv = ptable[recv_pid]
                .block_on(BlockedReason::Receiving(cap.phys(), cap.generation()))?;
            enqueue(&mut receivers, blocked_recv, ptable).map_err(|(e, proof)| {
                let _ = ptable.unblock_and_enqueue(recv_pid, proof);
                e
            })?;
            store_receivers(ep, &receivers);
            crate::sched::reset_effective(ptable, recv_pid);
            store_holder(ep, Some(recv_pid));
            Ok(IpcOutcome::Blocked)
        }
    }
}

pub fn do_try_recv(cap: &CapRef, recv_pid: Pid, ptable: &mut ProcessManager) -> IpcResult<Pid> {
    match dequeue_sender(cap, recv_pid, ptable)? {
        Some(blocked_sender) => handle_dequeued_sender(cap, recv_pid, blocked_sender, ptable),
        None => Err(crate::error::KernelError::WouldBlock),
    }
}

pub fn do_call(cap: &CapRef, caller_pid: Pid, ptable: &mut ProcessManager) -> CallResult {
    let caller_msg = ptable.exec(caller_pid).unwrap().ipc_message;

    let blocked_recv = {
        let mut pool = POOL.lock();
        let ep = pool.write_as::<EndpointObject>(cap.phys(), cap.generation())?;
        let mut receivers = load_receivers(ep);
        let dequeued = dequeue(&mut receivers, ptable);
        if let Some(br) = &dequeued {
            let recv_pid = br.pid();
            let old_holder = load_holder(ep);
            store_holder(ep, Some(recv_pid));
            match ptable[caller_pid]
                .block_on(BlockedReason::Calling(cap.phys(), cap.generation()))
            {
                Ok(blocked_caller) => {
                    if let Err((enq_err, caller_proof)) =
                        enqueue(&mut receivers, blocked_caller, ptable)
                    {
                        crate::show!(
                            ipc,
                            error,
                            "call enqueue caller pid {} failed {:?}",
                            caller_pid.raw(),
                            enq_err
                        );
                        {
                            let r = ptable.unblock_and_enqueue(caller_pid, caller_proof);
                            debug_assert!(r.is_ok());
                        }
                        let caller_exec = ptable.exec_mut(caller_pid).unwrap();
                        caller_exec.saved_context.rax = enq_err.to_errno() as u64;
                        caller_exec.seal_context();
                        let recv_proof = ptable[recv_pid].blocked_proof();
                        let re_enq = enqueue(&mut receivers, recv_proof, ptable);
                        debug_assert!(
                            re_enq.is_ok(),
                            "re-enqueue of just-dequeued receiver must succeed"
                        );
                        if let Err((_, recv_recovery)) = re_enq {
                            let recv_exec = ptable.exec_mut(recv_pid).unwrap();
                            recv_exec.saved_context.rax =
                                crate::error::KernelError::ResourceExhausted.to_errno() as u64;
                            recv_exec.seal_context();
                            {
                                let r = ptable.unblock_and_enqueue(recv_pid, recv_recovery);
                                debug_assert!(r.is_ok());
                            }
                        }
                        store_holder(ep, old_holder);
                        store_receivers(ep, &receivers);
                        return Err(enq_err);
                    }
                }
                Err(e) => {
                    let proof = ptable[recv_pid].blocked_proof();
                    let re_enq = enqueue(&mut receivers, proof, ptable);
                    debug_assert!(
                        re_enq.is_ok(),
                        "re-enqueue of just-dequeued receiver must succeed"
                    );
                    if let Err((enq_err, recv_recovery)) = re_enq {
                        crate::show!(
                            ipc,
                            error,
                            "call re-enqueue recv pid {} failed {:?}",
                            recv_pid.raw(),
                            enq_err
                        );
                        let recv_exec = ptable.exec_mut(recv_pid).unwrap();
                        recv_exec.saved_context.rax =
                            crate::error::KernelError::ResourceExhausted.to_errno() as u64;
                        recv_exec.seal_context();
                        {
                            let r = ptable.unblock_and_enqueue(recv_pid, recv_recovery);
                            debug_assert!(r.is_ok());
                        }
                    }
                    store_holder(ep, old_holder);
                    store_receivers(ep, &receivers);
                    return Err(e);
                }
            }
        }
        store_receivers(ep, &receivers);
        dequeued
    };

    match blocked_recv {
        Some(blocked_recv) => {
            let recv_pid = blocked_recv.pid();
            let caller_prio = ptable[caller_pid].effective_priority();

            ptable.unblock_and_enqueue(recv_pid, blocked_recv)?;
            let recv_exec = ptable.exec_mut(recv_pid).unwrap();
            recv_exec.ipc_message = caller_msg;
            recv_exec.ipc_badge = caller_pid.raw() as u64;
            message::inject_into_context(&mut recv_exec.saved_context, &caller_msg);
            recv_exec.saved_context.rax = caller_pid.raw() as u64;
            recv_exec.reply_target = Some(caller_pid);
            recv_exec.seal_context();

            crate::sched::boost_effective(ptable, recv_pid, caller_prio);

            Ok(AlwaysBlocked)
        }
        None => {
            let mut pool = POOL.lock();
            let (holder, donor_prio) = {
                let ep = pool.write_as::<EndpointObject>(cap.phys(), cap.generation())?;
                let mut senders = load_senders(ep);
                let blocked_caller = ptable[caller_pid]
                    .block_on(BlockedReason::Calling(cap.phys(), cap.generation()))?;
                enqueue(&mut senders, blocked_caller, ptable).map_err(|(e, proof)| {
                    let _ = ptable.unblock_and_enqueue(caller_pid, proof);
                    e
                })?;
                store_senders(ep, &senders);
                let holder = load_holder(ep);
                (holder, ptable[caller_pid].effective_priority())
            };
            if let Some(h) = holder {
                crate::sched::propagate_priority(ptable, &pool, h, donor_prio);
            }

            Ok(AlwaysBlocked)
        }
    }
}

pub fn remove_from_recv(ep: &mut EndpointObject, pid: Pid, ptable: &mut ProcessManager) {
    let mut receivers = load_receivers(ep);
    remove_pid(&mut receivers, pid, ptable);
    store_receivers(ep, &receivers);
}

fn unblock_callers_in_receivers(queue: &mut PidQueue, ptable: &mut ProcessManager) {
    let mut new_head: Option<Pid> = None;
    let mut new_tail: Option<Pid> = None;
    let mut new_len: u16 = 0;
    let mut cursor = queue.head;
    let mut steps = 0u32;

    core::iter::from_fn(|| {
        cursor.filter(|_| steps < MAX_PIDS as u32).inspect(|&pid| {
            steps += 1;
            let next = ptable[pid].next_ipc;
            ptable[pid].next_ipc = None;
            cursor = next;

            if matches!(
                ptable[pid].blocked_reason(),
                Some(BlockedReason::Calling(_, _))
            ) {
                let proof = ptable[pid].blocked_proof();
                match ptable.unblock_and_enqueue(pid, proof) {
                    Ok(()) => {
                        let exec = ptable.exec_mut(pid).unwrap();
                        exec.saved_context.rax =
                            crate::error::KernelError::InvalidObject.to_errno() as u64;
                        exec.seal_context();
                    }
                    Err(e) => {
                        crate::show!(
                            ipc,
                            error,
                            "unblock caller pid {} on holder death failed {:?}",
                            pid.raw(),
                            e
                        );
                    }
                }
            } else {
                if let Some(tail) = new_tail {
                    ptable[tail].next_ipc = Some(pid);
                } else {
                    new_head = Some(pid);
                }
                new_tail = Some(pid);
                new_len = new_len.saturating_add(1);
            }
        })
    })
    .count();

    *queue = PidQueue::from_parts(new_head, new_tail, new_len);

    assert_queue_valid(queue, ptable);
}

pub fn remove_from_queues(
    pid: Pid,
    pool: &mut crate::cap::pool::ObjectPool,
    ptable: &mut ProcessManager,
) {
    let blocked = ptable.get(pid).and_then(|s| s.blocked_reason());

    match blocked {
        Some(BlockedReason::Sending(phys, generation)) | Some(BlockedReason::Calling(phys, generation)) => {
            if let Ok(ep) = pool.write_as::<EndpointObject>(phys, generation) {
                let mut senders = load_senders(ep);
                remove_pid(&mut senders, pid, ptable);
                store_senders(ep, &senders);

                let mut receivers = load_receivers(ep);
                let holder = load_holder(ep);
                match holder {
                    Some(h) if h == pid => {
                        store_holder(ep, None);
                        unblock_callers_in_receivers(&mut receivers, ptable);
                        store_receivers(ep, &receivers);
                    }
                    _ => {
                        store_receivers(ep, &receivers);
                    }
                }
            }
        }
        Some(BlockedReason::Receiving(phys, generation)) => {
            if let Ok(ep) = pool.write_as::<EndpointObject>(phys, generation) {
                let mut receivers = load_receivers(ep);
                remove_pid(&mut receivers, pid, ptable);
                store_receivers(ep, &receivers);

                let holder = load_holder(ep);
                match holder {
                    Some(h) if h == pid => {
                        store_holder(ep, None);
                        let mut receivers = load_receivers(ep);
                        unblock_callers_in_receivers(&mut receivers, ptable);
                        store_receivers(ep, &receivers);
                    }
                    _ => {}
                }
            }
        }
        Some(BlockedReason::WaitingNotification(phys, generation)) => {
            if let Ok(notif) = pool
                .write_as::<lancer_core::object_layout::NotificationObject>(phys, generation)
            {
                let raw = pid.raw();
                let old_waiters = notif.waiters;
                let retained_count = (0..notif.waiter_count as usize)
                    .filter(|&i| old_waiters[i] != raw)
                    .fold(0u8, |dst_idx, i| {
                        notif.waiters[dst_idx as usize] = old_waiters[i];
                        dst_idx + 1
                    });
                (retained_count as usize..notif.waiters.len()).for_each(|i| {
                    notif.waiters[i] = NONE_SENTINEL;
                });
                notif.waiter_count = retained_count;
            }
        }
        None => {}
    }
}