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::{IpcMessage, Reply};
use crate::proc::{BlockedReason, PROCESSES, ProcessState};
use lancer_core::object_layout::EndpointObject;

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

        let client_created = ptable.allocate(&mut allocator).expect("alloc client");
        let server_created = ptable.allocate(&mut allocator).expect("alloc server");
        let sender_created = ptable.allocate(&mut allocator).expect("alloc sender");
        ptable.start(client_created).expect("start client");
        ptable.start(server_created).expect("start server");
        ptable.start(sender_created).expect("start sender");
        let client_pid = client_created.pid();
        let server_pid = server_created.pid();
        let sender_pid = sender_created.pid();

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

        let cap = CapRef::new(ObjectTag::Endpoint, ep_id, Rights::ALL, ep_gen);

        ptable.exec_mut(client_pid).unwrap().ipc_message =
            IpcMessage::from_regs([0x1111, 0, 0, 0, 0, 0]);
        ptable.simulate_dispatch(client_pid);
        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 for store");
            endpoint::store_receivers(ep, &receivers);
        }

        let call_result = endpoint::do_call(&cap, client_pid, &mut ptable);
        match call_result {
            Ok(AlwaysBlocked) => {}
            _ => panic!("expected Blocked from do_call"),
        }

        assert!(
            ptable.exec(server_pid).unwrap().reply_target == Some(client_pid),
            "server should have reply_target set to client"
        );

        ptable.exec_mut(sender_pid).unwrap().ipc_message =
            IpcMessage::from_regs([0x2222, 0, 0, 0, 0, 0]);
        ptable.simulate_dispatch(sender_pid);
        let sender_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, sender_blocked, &mut ptable).expect("enqueue sender");
            let ep = pool
                .write_as::<EndpointObject>(ep_id, ep_gen)
                .expect("get ep for store");
            endpoint::store_senders(ep, &senders);
        }

        let full_reply = IpcMessage::<Reply>::from_reply_regs([0xAAAA, 0, 0, 0, 0]).into_full();
        {
            let client_exec = ptable.exec_mut(client_pid).unwrap();
            let proof = unsafe { crate::types::BlockedPid::trust(client_pid) };
            client_exec.ipc_message = full_reply;
            message::inject_into_context(&mut client_exec.saved_context, &full_reply);
            client_exec.saved_context.rax = server_pid.raw() as u64;
            ptable.unblock_and_enqueue(client_pid, proof).expect("unblock client");
        }
        ptable.exec_mut(server_pid).unwrap().reply_target = None;

        let recv_result = endpoint::do_recv(&cap, server_pid, &mut ptable);
        match recv_result {
            Ok(IpcOutcome::Done(s)) => {
                assert!(s == sender_pid, "server should receive from sender");
            }
            _ => panic!("expected Done from do_recv after reply"),
        }

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

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

        let server_created = ptable.allocate(&mut allocator).expect("alloc server");
        let receiver_created = ptable.allocate(&mut allocator).expect("alloc receiver");
        ptable.start(server_created).expect("start server");
        ptable.start(receiver_created).expect("start receiver");
        let server_pid = server_created.pid();
        let receiver_pid = receiver_created.pid();

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

        ptable.simulate_dispatch(receiver_pid);
        let blocked_recv = ptable[receiver_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 for store");
            endpoint::store_receivers(ep, &receivers);
        }

        ptable.exec_mut(server_pid).unwrap().reply_target = Some(receiver_pid);

        let is_call_target = ptable[receiver_pid].state() == ProcessState::Blocked
            && matches!(
                ptable[receiver_pid].blocked_reason(),
                Some(BlockedReason::Calling(id, _)) if id == ep_id
            );

        assert!(
            !is_call_target,
            "a Receiving process should NOT be accepted as a Calling reply target"
        );

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

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

        let server_created = ptable.allocate(&mut allocator).expect("alloc server");
        ptable.start(server_created).expect("start server");
        let server_pid = server_created.pid();

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

        ptable.exec_mut(server_pid).unwrap().reply_target = None;

        let reply_pid = crate::types::Pid::try_new(u32::MAX);
        assert!(reply_pid.is_none(), "Pid::try_new(MAX) should return None");

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