use crate::cap::object::ObjectTag;
use crate::cap::pool::POOL;
use crate::cap::table::{CapRef, Rights};
use crate::ipc::{AlwaysBlocked, IpcOutcome, endpoint, message};
use crate::proc::context::{CpuContext, IpcMessage};
use crate::proc::{BlockedReason, PROCESSES, ProcessState};
use lancer_core::object_layout::EndpointObject;

crate::kernel_test!(
    fn recv_fast_path_delivers_message() {
        let mut allocator = crate::mem::phys::BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();

        let sender_created = ptable.allocate(&mut allocator).expect("alloc sender");
        let recv_created = ptable.allocate(&mut allocator).expect("alloc receiver");
        ptable.start(sender_created).expect("start sender");
        ptable.start(recv_created).expect("start receiver");
        let sender_pid = sender_created.pid();
        let recv_pid = recv_created.pid();

        let (ep_id, ep_gen) =
            crate::tests::helpers::alloc_endpoint(&mut POOL.lock()).expect("alloc endpoint");

        let msg = IpcMessage::from_regs([
            0xCAFE_0001,
            0xCAFE_0002,
            0xCAFE_0003,
            0xCAFE_0004,
            0xCAFE_0005,
            0xCAFE_0006,
        ]);
        ptable.exec_mut(sender_pid).unwrap().ipc_message = msg;

        ptable.simulate_dispatch(sender_pid);
        let blocked = ptable[sender_pid]
            .block_on(BlockedReason::Sending(ep_id, ep_gen))
            .expect("block sender");

        {
            let mut pool = POOL.lock();
            let ep = pool
                .write_as::<EndpointObject>(ep_id, ep_gen)
                .expect("get ep");
            let mut senders = endpoint::load_senders(ep);
            endpoint::enqueue(&mut senders, blocked, &mut ptable).expect("enqueue sender");
            let ep = pool
                .write_as::<EndpointObject>(ep_id, ep_gen)
                .expect("get ep");
            endpoint::store_senders(ep, &senders);
        }

        let cap = CapRef::new(ObjectTag::Endpoint, ep_id, Rights::ALL, ep_gen);
        let result = endpoint::do_recv(&cap, recv_pid, &mut ptable);

        let returned_sender = match result {
            Ok(IpcOutcome::Done(pid)) => pid,
            _ => panic!("expected Done from do_recv fast path"),
        };
        assert!(
            returned_sender == sender_pid,
            "do_recv should return the sender's pid"
        );

        let recv_msg = ptable.exec(recv_pid).unwrap().ipc_message;

        let mut ctx = CpuContext::zero();
        message::inject_into_context(&mut ctx, &recv_msg);
        ctx.rax = returned_sender.raw() as u64;

        assert!(
            ctx.rsi == 0xCAFE_0001,
            "msg[0] must be delivered to ctx.rsi"
        );
        assert!(
            ctx.rdx == 0xCAFE_0002,
            "msg[1] must be delivered to ctx.rdx"
        );
        assert!(
            ctx.r10 == 0xCAFE_0003,
            "msg[2] must be delivered to ctx.r10"
        );
        assert!(ctx.r8 == 0xCAFE_0004, "msg[3] must be delivered to ctx.r8");
        assert!(ctx.r9 == 0xCAFE_0005, "msg[4] must be delivered to ctx.r9");
        assert!(
            ctx.r12 == 0xCAFE_0006,
            "msg[5] must be delivered to ctx.r12"
        );

        let _ = POOL.lock().dec_ref_phys(ep_id, ep_gen);
        ptable.destroy(sender_pid, &mut allocator);
        ptable.destroy(recv_pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn call_fast_path_sets_reply_target() {
        let mut allocator = crate::mem::phys::BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();

        let caller_created = ptable.allocate(&mut allocator).expect("alloc caller");
        let recv_created = ptable.allocate(&mut allocator).expect("alloc receiver");
        ptable.start(caller_created).expect("start caller");
        ptable.start(recv_created).expect("start receiver");
        let caller_pid = caller_created.pid();
        let recv_pid = recv_created.pid();

        let (ep_id, ep_gen) =
            crate::tests::helpers::alloc_endpoint(&mut POOL.lock()).expect("alloc endpoint");

        ptable.simulate_dispatch(recv_pid);
        let blocked_recv = ptable[recv_pid]
            .block_on(BlockedReason::Receiving(ep_id, ep_gen))
            .expect("block receiver");

        {
            let mut pool = POOL.lock();
            let ep = pool
                .write_as::<EndpointObject>(ep_id, ep_gen)
                .expect("get ep");
            let mut receivers = endpoint::load_receivers(ep);
            endpoint::enqueue(&mut receivers, blocked_recv, &mut ptable).expect("enqueue receiver");
            let ep = pool
                .write_as::<EndpointObject>(ep_id, ep_gen)
                .expect("get ep");
            endpoint::store_receivers(ep, &receivers);
        }

        let msg = IpcMessage::from_regs([0x1234, 0x5678, 0, 0, 0, 0]);
        ptable.exec_mut(caller_pid).unwrap().ipc_message = msg;
        ptable.simulate_dispatch(caller_pid);

        let cap = CapRef::new(ObjectTag::Endpoint, ep_id, Rights::ALL, ep_gen);
        let result = endpoint::do_call(&cap, caller_pid, &mut ptable);

        match result {
            Ok(AlwaysBlocked) => {}
            _ => panic!("expected Blocked from do_call fast path"),
        }

        assert!(
            ptable[recv_pid].state() == ProcessState::Ready,
            "receiver should be unblocked after do_call fast path"
        );

        assert!(
            ptable.exec(recv_pid).unwrap().saved_context.rax == caller_pid.raw() as u64,
            "receiver's rax should hold the caller's pid"
        );

        assert!(
            ptable.exec(recv_pid).unwrap().reply_target == Some(caller_pid),
            "do_call fast path must set reply_target so receiver can reply"
        );

        assert!(
            ptable[caller_pid].state() == ProcessState::Blocked,
            "caller should be blocked waiting for reply"
        );

        assert!(
            matches!(
                ptable[caller_pid].blocked_reason(),
                Some(BlockedReason::Calling(id, _)) if id == ep_id
            ),
            "caller should be blocked as Calling on the endpoint"
        );

        let _ = POOL.lock().dec_ref_phys(ep_id, ep_gen);
        ptable.destroy(caller_pid, &mut allocator);
        ptable.destroy(recv_pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn send_skips_call_waiters() {
        let mut allocator = crate::mem::phys::BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();

        let server_created = ptable.allocate(&mut allocator).expect("alloc server");
        let client_a_created = ptable.allocate(&mut allocator).expect("alloc client_a");
        let client_b_created = ptable.allocate(&mut allocator).expect("alloc client_b");
        ptable.start(server_created).expect("start server");
        ptable.start(client_a_created).expect("start client_a");
        ptable.start(client_b_created).expect("start client_b");
        let server_pid = server_created.pid();
        let client_a_pid = client_a_created.pid();
        let client_b_pid = client_b_created.pid();

        let (ep_id, ep_gen) =
            crate::tests::helpers::alloc_endpoint(&mut POOL.lock()).expect("alloc endpoint");

        ptable.simulate_dispatch(server_pid);
        let blocked_server = ptable[server_pid]
            .block_on(BlockedReason::Receiving(ep_id, ep_gen))
            .expect("block server");

        {
            let mut pool = POOL.lock();
            let ep = pool
                .write_as::<EndpointObject>(ep_id, ep_gen)
                .expect("get ep");
            let mut receivers = endpoint::load_receivers(ep);
            endpoint::enqueue(&mut receivers, blocked_server, &mut ptable).expect("enqueue server");
            let ep = pool
                .write_as::<EndpointObject>(ep_id, ep_gen)
                .expect("get ep");
            endpoint::store_receivers(ep, &receivers);
        }

        let msg_a = IpcMessage::from_regs([0xAAAA, 0, 0, 0, 0, 0]);
        ptable.exec_mut(client_a_pid).unwrap().ipc_message = msg_a;
        ptable.simulate_dispatch(client_a_pid);

        let cap = CapRef::new(ObjectTag::Endpoint, ep_id, Rights::ALL, ep_gen);
        let call_result = endpoint::do_call(&cap, client_a_pid, &mut ptable);
        match call_result {
            Ok(AlwaysBlocked) => {}
            _ => panic!("expected Blocked from client_a do_call"),
        }

        assert!(
            ptable[server_pid].state() == ProcessState::Ready,
            "server should be Ready after receiving call"
        );
        assert!(
            ptable[client_a_pid].state() == ProcessState::Blocked,
            "client_a should be Blocked waiting for reply"
        );

        let msg_b = IpcMessage::from_regs([0xBBBB, 0, 0, 0, 0, 0]);
        ptable.exec_mut(client_b_pid).unwrap().ipc_message = msg_b;
        ptable.simulate_dispatch(client_b_pid);

        let send_result = endpoint::do_send(&cap, client_b_pid, &mut ptable);
        match send_result {
            Ok(IpcOutcome::Blocked) => {}
            _ => panic!("expected Blocked from client_b do_send (no genuine receiver)"),
        }

        assert!(
            ptable[client_a_pid].state() == ProcessState::Blocked,
            "client_a must still be blocked. send must not wake call-waiters"
        );
        assert!(
            matches!(
                ptable[client_a_pid].blocked_reason(),
                Some(BlockedReason::Calling(id, _)) if id == ep_id
            ),
            "client_a should still have Calling blocked reason"
        );

        assert!(
            ptable[client_b_pid].state() == ProcessState::Blocked,
            "client_b should be blocked as sender (no genuine receiver available)"
        );
        assert!(
            matches!(
                ptable[client_b_pid].blocked_reason(),
                Some(BlockedReason::Sending(id, _)) if id == ep_id
            ),
            "client_b should have Sending blocked reason"
        );

        let _ = POOL.lock().dec_ref_phys(ep_id, ep_gen);
        ptable.destroy(server_pid, &mut allocator);
        ptable.destroy(client_a_pid, &mut allocator);
        ptable.destroy(client_b_pid, &mut allocator);
    }
);