//! Actual cooperative EL1 stack switches for a full-channel sender and an
//! empty-channel receiver. IRQ stays masked; no timer/deadline claim here.
use super::*;
use turkos_kernel::ipc::{self, ChannelTable, ReceiveOutcome, SendOutcome};
use turkos_kernel::task::{self, Scheduler, State, TaskId, WaitKind, WakeCause};

const STACK_BYTES: usize = 16 * 1024;
const CANARY_A: u64 = 0x4950_4341_5354_414b;
const CANARY_B: u64 = 0x4950_4342_5354_414b;
#[repr(C, align(4096))]
struct TaskStack([u8; STACK_BYTES]);

static mut A_STACK: TaskStack = TaskStack([0; STACK_BYTES]);
static mut B_STACK: TaskStack = TaskStack([0; STACK_BYTES]);
static mut A_CONTEXT: task::CpuContext = task::CpuContext::EMPTY;
static mut B_CONTEXT: task::CpuContext = task::CpuContext::EMPTY;
static mut BOOT_CONTEXT: task::CpuContext = task::CpuContext::EMPTY;
static mut SCHEDULER: Scheduler = Scheduler::new();
static mut CHANNELS: ChannelTable = ChannelTable::new();
static mut IDS: [TaskId; 2] = [TaskId {
    slot: 0,
    generation: 0,
}; 2];
static mut CHANNEL_ID: u64 = 0;
static mut SECOND_CHANNEL_ID: u64 = 0;
static mut TIMEOUT_CHANNEL_ID: u64 = 0;
static mut TIMEOUT_DEADLINE: u64 = 0;
static mut A_STEPS: u64 = 0;

unsafe extern "C" {
    fn __turkos_swap_task_context(old: *mut task::CpuContext, new: *const task::CpuContext);
}

fn fail(reason: &str) -> ! {
    write_str("TURKOS/IPC-WAIT-ERROR: ");
    write_line(reason);
    park()
}

fn id(index: usize) -> TaskId {
    unsafe { read_volatile(core::ptr::addr_of!(IDS))[index] }
}

fn channel() -> u64 {
    unsafe { read_volatile(core::ptr::addr_of!(CHANNEL_ID)) }
}

fn second_channel() -> u64 {
    unsafe { read_volatile(core::ptr::addr_of!(SECOND_CHANNEL_ID)) }
}

fn timeout_channel() -> u64 {
    unsafe { read_volatile(core::ptr::addr_of!(TIMEOUT_CHANNEL_ID)) }
}

fn counter() -> u64 {
    let value: u64;
    unsafe {
        asm!("isb", "mrs {value}, CNTPCT_EL0", value = out(reg) value, options(nomem, nostack, preserves_flags))
    };
    value
}

fn counter_frequency() -> u64 {
    let value: u64;
    unsafe {
        asm!("mrs {value}, CNTFRQ_EL0", value = out(reg) value, options(nomem, nostack, preserves_flags))
    };
    value
}

fn scheduler_ptr() -> *mut Scheduler {
    core::ptr::addr_of_mut!(SCHEDULER)
}

fn context(index: usize) -> *mut task::CpuContext {
    if index == 0 {
        core::ptr::addr_of_mut!(A_CONTEXT)
    } else {
        core::ptr::addr_of_mut!(B_CONTEXT)
    }
}

fn switch(old: usize, new: usize) {
    unsafe { __turkos_swap_task_context(context(old), context(new)) };
}

extern "C" fn task_a() -> ! {
    let a = id(0);
    let b = id(1);
    let channel = channel();
    for byte in [b'0', b'1', b'2', b'3'] {
        let result = unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).send_or_block(
                channel,
                a,
                &[byte],
                &mut *core::ptr::addr_of_mut!(SCHEDULER),
                0,
            )
        };
        if result != Ok(SendOutcome::Sent) {
            fail("initial four sends");
        }
    }
    let mut pending = [b'4'];
    let result = unsafe {
        (&mut *core::ptr::addr_of_mut!(CHANNELS)).send_or_block(
            channel,
            a,
            &pending,
            &mut *core::ptr::addr_of_mut!(SCHEDULER),
            0,
        )
    };
    if result != Ok(SendOutcome::Blocked(Some(b))) {
        fail("fifth send must block and select B");
    }
    pending[0] = b'X';
    write_line("TURKOS/IPC-WAIT-FULL: QUEUED=0x0000000000000004 A=BLOCKED B=RUNNING");
    switch(0, 1);

    if unsafe { (&*scheduler_ptr()).running() } != Some(a)
        || unsafe { (&mut *scheduler_ptr()).take_wake_cause(a) }
            != Ok(Some(WakeCause::SpaceAvailable))
        || pending[0] != b'X'
        || unsafe { read_volatile(core::ptr::addr_of!(A_STEPS)) } != 0
    {
        fail("A resumed only after B released queue space");
    }
    unsafe { write_volatile(core::ptr::addr_of_mut!(A_STEPS), 1) };
    write_line("TURKOS/IPC-WAIT-A-WAKE: CAUSE=SPACE STEPS=0x0000000000000001");
    if unsafe {
        (&mut *core::ptr::addr_of_mut!(CHANNELS)).send_or_block(
            channel,
            a,
            b"PING",
            &mut *scheduler_ptr(),
            0,
        )
    } != Ok(SendOutcome::Sent)
        || unsafe { (&*scheduler_ptr()).state(b) } != Ok(State::Ready)
    {
        fail("A send must wake B");
    }
    write_line("TURKOS/IPC-WAIT-SEND: VALUE=PING B=READY");
    if unsafe { (&mut *scheduler_ptr()).yield_current(false) } != Ok(Some(b)) {
        fail("A yield to B");
    }
    switch(0, 1);

    if unsafe { (&*scheduler_ptr()).running() } != Some(a) {
        fail("A must resume for close-wake setup");
    }
    for byte in [b'0', b'1', b'2', b'3'] {
        if unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).send_or_block(
                channel,
                a,
                &[byte],
                &mut *scheduler_ptr(),
                0,
            )
        } != Ok(SendOutcome::Sent)
        {
            fail("close scenario fill");
        }
    }
    if unsafe {
        (&mut *core::ptr::addr_of_mut!(CHANNELS)).send_or_block(
            channel,
            a,
            b"LATE",
            &mut *scheduler_ptr(),
            0,
        )
    } != Ok(SendOutcome::Blocked(Some(b)))
    {
        fail("close scenario must block A");
    }
    write_line("TURKOS/IPC-WAIT-CLOSE-FULL: A=BLOCKED B=RUNNING QUEUED=0x0000000000000004");
    switch(0, 1);

    if unsafe { (&*scheduler_ptr()).running() } != Some(a)
        || unsafe { (&mut *scheduler_ptr()).take_wake_cause(a) }
            != Ok(Some(WakeCause::ChannelClosed))
        || unsafe { (&*scheduler_ptr()).ready_len() } != 1
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).send_or_block(
                channel,
                a,
                b"AGAIN",
                &mut *scheduler_ptr(),
                0,
            )
        } != Err(ipc::Error::NoReceiver)
    {
        fail("closed sender must wake once and reject retry");
    }
    write_line("TURKOS/IPC-WAIT-CLOSE-A-WAKE: CAUSE=CLOSED RETRY=NO-RECEIVER READY=1");
    let second = unsafe { (&mut *core::ptr::addr_of_mut!(CHANNELS)).create(a, b) }
        .unwrap_or_else(|_| fail("second channel create"));
    unsafe { write_volatile(core::ptr::addr_of_mut!(SECOND_CHANNEL_ID), second) };
    if second != channel + 1
        || unsafe { (&mut *scheduler_ptr()).yield_current(false) } != Ok(Some(b))
    {
        fail("A yield to second-channel receiver");
    }
    switch(0, 1);

    if unsafe { (&*scheduler_ptr()).running() } != Some(a)
        || unsafe { (&*scheduler_ptr()).state(b) }
            != Ok(State::Blocked {
                channel: second,
                kind: WaitKind::Receive,
                deadline: None,
            })
        || unsafe { (&*scheduler_ptr()).ready_len() } != 0
        || unsafe { (&mut *core::ptr::addr_of_mut!(CHANNELS)).destroy(second, a) }
            != Err(ipc::Error::RequiresScheduler)
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).destroy_wait(second, a, &mut *scheduler_ptr())
        } != Ok(())
        || unsafe { (&*scheduler_ptr()).state(b) } != Ok(State::Ready)
    {
        fail("destroy must wake blocked B exactly once");
    }
    write_line("TURKOS/IPC-WAIT-DESTROY-RECEIVER: CHANNEL=0x0000000000000002 B=READY READY=1");
    if unsafe { (&mut *scheduler_ptr()).yield_current(false) } != Ok(Some(b)) {
        fail("A yield to destroyed-channel receiver");
    }
    switch(0, 1);

    let timeout_channel = timeout_channel();
    let deadline = unsafe { read_volatile(core::ptr::addr_of!(TIMEOUT_DEADLINE)) };
    let early = counter();
    if timeout_channel != second + 1
        || early >= deadline
        || unsafe { (&*scheduler_ptr()).state(b) }
            != Ok(State::Blocked {
                channel: timeout_channel,
                kind: WaitKind::Receive,
                deadline: Some(deadline),
            })
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).poll_deadlines(&mut *scheduler_ptr(), early)
        } != Ok(0)
        || unsafe { (&*scheduler_ptr()).ready_len() } != 0
    {
        fail("deadline must leave B blocked before physical counter threshold");
    }
    write_str("TURKOS/IPC-WAIT-TIMEOUT-BEFORE: B=BLOCKED READY=0 NOW=0x");
    write_hex(early as usize);
    write_str(" DEADLINE=0x");
    write_hex(deadline as usize);
    write_str("\r\n");
    // A is deliberately the active clock service in this narrow experiment.
    // This is not timer-IRQ-driven idle sleep; B itself consumes no cycles.
    while counter() < deadline {
        core::hint::spin_loop();
    }
    let due = counter();
    if unsafe {
        (&mut *core::ptr::addr_of_mut!(CHANNELS)).poll_deadlines(&mut *scheduler_ptr(), due)
    } != Ok(1)
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).poll_deadlines(&mut *scheduler_ptr(), due)
        } != Ok(0)
        || unsafe { (&*scheduler_ptr()).state(b) } != Ok(State::Ready)
        || unsafe { (&*scheduler_ptr()).ready_len() } != 1
    {
        fail("physical deadline must wake B exactly once");
    }
    write_str("TURKOS/IPC-WAIT-TIMEOUT-DUE: B=READY WAKE=1 REPEAT=0 READY=1 NOW=0x");
    write_hex(due as usize);
    write_str(" DEADLINE=0x");
    write_hex(deadline as usize);
    write_str("\r\n");
    if unsafe { (&mut *scheduler_ptr()).yield_current(false) } != Ok(Some(b)) {
        fail("A yield to timed-out B");
    }
    switch(0, 1);

    if unsafe { (&*scheduler_ptr()).running() } != Some(a)
        || unsafe { (&*scheduler_ptr()).state(b) } != Ok(State::Exited { code: 0x0b })
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).destroy_wait(
                timeout_channel,
                a,
                &mut *scheduler_ptr(),
            )
        } != Ok(())
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).destroy_wait(
                channel,
                a,
                &mut *scheduler_ptr(),
            )
        } != Ok(())
        || unsafe { (&*core::ptr::addr_of!(CHANNELS)).queued(channel) }
            != Err(ipc::Error::InvalidChannel)
    {
        fail("A final state after B exit");
    }
    write_line("TURKOS/IPC-WAIT-A-EXIT: B=EXITED QUEUED=0x0000000000000000");
    if unsafe { (&mut *scheduler_ptr()).exit_current(0x0a) } != Ok(None) {
        fail("A exit to boot");
    }
    unsafe { __turkos_swap_task_context(context(0), core::ptr::addr_of!(BOOT_CONTEXT)) };
    fail("A returned after exit")
}

extern "C" fn task_b() -> ! {
    let a = id(0);
    let b = id(1);
    let channel = channel();
    if unsafe { (&*scheduler_ptr()).running() } != Some(b)
        || unsafe { (&*scheduler_ptr()).state(a) }
            != Ok(State::Blocked {
                channel,
                kind: WaitKind::Send,
                deadline: None,
            })
        || unsafe { (&*scheduler_ptr()).ready_len() } != 0
        || unsafe { read_volatile(core::ptr::addr_of!(A_STEPS)) } != 0
    {
        fail("B entered after true A block");
    }
    for expected in [b'0', b'1', b'2', b'3', b'4'] {
        let outcome = unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).receive_or_block(
                channel,
                b,
                &mut *scheduler_ptr(),
                None,
                0,
            )
        }
        .unwrap_or_else(|_| fail("B drain receive"));
        if outcome.message_bytes() != Some(&[expected][..]) {
            fail("B FIFO or pending send copy");
        }
    }
    if unsafe { (&*scheduler_ptr()).state(a) } != Ok(State::Ready)
        || unsafe { (&*core::ptr::addr_of!(CHANNELS)).queued(channel) } != Ok(0)
        || unsafe { read_volatile(core::ptr::addr_of!(A_STEPS)) } != 0
    {
        fail("A cannot run during B drain");
    }
    write_line("TURKOS/IPC-WAIT-DRAIN: ORDER=01234 A=READY A-STEPS=0x0000000000000000");
    let result = unsafe {
        (&mut *core::ptr::addr_of_mut!(CHANNELS)).receive_or_block(
            channel,
            b,
            &mut *scheduler_ptr(),
            None,
            0,
        )
    };
    if result != Ok(ReceiveOutcome::Blocked(Some(a)))
        || unsafe { (&*scheduler_ptr()).ready_len() } != 0
        || unsafe { (&*scheduler_ptr()).state(b) }
            != Ok(State::Blocked {
                channel,
                kind: WaitKind::Receive,
                deadline: None,
            })
    {
        fail("empty receive must block B and select A");
    }
    write_line("TURKOS/IPC-WAIT-EMPTY: QUEUED=0x0000000000000000 B=BLOCKED A=RUNNING");
    switch(1, 0);

    if unsafe { (&*scheduler_ptr()).running() } != Some(b)
        || unsafe { (&mut *scheduler_ptr()).take_wake_cause(b) }
            != Ok(Some(WakeCause::MessageAvailable))
    {
        fail("B resumed by A message");
    }
    let ping = unsafe {
        (&mut *core::ptr::addr_of_mut!(CHANNELS)).receive_or_block(
            channel,
            b,
            &mut *scheduler_ptr(),
            None,
            0,
        )
    }
    .unwrap_or_else(|_| fail("B PING receive"));
    if ping.message_bytes() != Some(&b"PING"[..]) {
        fail("B wake payload");
    }
    write_line("TURKOS/IPC-WAIT-B-WAKE: CAUSE=MESSAGE VALUE=PING");
    if unsafe { (&mut *scheduler_ptr()).yield_current(false) } != Ok(Some(a)) {
        fail("B yield to close scenario A");
    }
    switch(1, 0);
    if unsafe { (&*scheduler_ptr()).running() } != Some(b)
        || unsafe { (&*scheduler_ptr()).state(a) }
            != Ok(State::Blocked {
                channel,
                kind: WaitKind::Send,
                deadline: None,
            })
        || unsafe { (&*scheduler_ptr()).ready_len() } != 0
        || unsafe { (&mut *core::ptr::addr_of_mut!(CHANNELS)).close_receiver(channel, b) }
            != Err(ipc::Error::RequiresScheduler)
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).close_receiver_wait(
                channel,
                b,
                &mut *scheduler_ptr(),
            )
        } != Ok(())
        || unsafe { (&*scheduler_ptr()).state(a) } != Ok(State::Ready)
    {
        fail("close must wake blocked A exactly once");
    }
    write_line("TURKOS/IPC-WAIT-CLOSE-SENDER: A=READY CAUSE=CLOSED READY=1");
    if unsafe { (&mut *scheduler_ptr()).yield_current(false) } != Ok(Some(a)) {
        fail("B yield to closed sender");
    }
    switch(1, 0);

    let second = second_channel();
    if second != channel + 1
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).receive_or_block(
                second,
                b,
                &mut *scheduler_ptr(),
                None,
                0,
            )
        } != Ok(ReceiveOutcome::Blocked(Some(a)))
        || unsafe { (&*scheduler_ptr()).ready_len() } != 0
    {
        fail("B block on second empty channel");
    }
    write_line("TURKOS/IPC-WAIT-CLOSE-EMPTY: B=BLOCKED A=RUNNING QUEUED=0x0000000000000000");
    switch(1, 0);
    if unsafe { (&*scheduler_ptr()).running() } != Some(b)
        || unsafe { (&mut *scheduler_ptr()).take_wake_cause(b) }
            != Ok(Some(WakeCause::ChannelClosed))
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).receive_or_block(
                second,
                b,
                &mut *scheduler_ptr(),
                None,
                0,
            )
        } != Err(ipc::Error::InvalidChannel)
    {
        fail("destroyed receiver must wake and reject stale channel");
    }
    write_line("TURKOS/IPC-WAIT-DESTROY-B-WAKE: CAUSE=CLOSED RETRY=INVALID-CHANNEL");
    let third = unsafe { (&mut *core::ptr::addr_of_mut!(CHANNELS)).create(a, b) }
        .unwrap_or_else(|_| fail("timeout channel create"));
    let frequency = counter_frequency();
    let now = counter();
    let duration = frequency / 2;
    let deadline = now
        .checked_add(duration)
        .unwrap_or_else(|| fail("timeout overflow"));
    if third != second + 1 || duration == 0 {
        fail("physical clock or third channel invalid");
    }
    unsafe {
        write_volatile(core::ptr::addr_of_mut!(TIMEOUT_CHANNEL_ID), third);
        write_volatile(core::ptr::addr_of_mut!(TIMEOUT_DEADLINE), deadline);
    }
    if unsafe {
        (&mut *core::ptr::addr_of_mut!(CHANNELS)).receive_or_block(
            third,
            b,
            &mut *scheduler_ptr(),
            Some(deadline),
            now,
        )
    } != Ok(ReceiveOutcome::Blocked(Some(a)))
    {
        fail("timed receive must block B");
    }
    write_str("TURKOS/IPC-WAIT-TIMEOUT-BLOCK: CHANNEL=0x");
    write_hex(third as usize);
    write_str(" B=BLOCKED A=RUNNING START=0x");
    write_hex(now as usize);
    write_str(" DEADLINE=0x");
    write_hex(deadline as usize);
    write_str(" FREQ=0x");
    write_hex(frequency as usize);
    write_str("\r\n");
    switch(1, 0);
    if unsafe { (&*scheduler_ptr()).running() } != Some(b)
        || unsafe { (&mut *scheduler_ptr()).take_wake_cause(b) } != Ok(Some(WakeCause::TimedOut))
        || counter() < deadline
        || unsafe {
            (&mut *core::ptr::addr_of_mut!(CHANNELS)).receive_or_block(
                third,
                b,
                &mut *scheduler_ptr(),
                Some(deadline),
                counter(),
            )
        } != Err(ipc::Error::TimedOut)
    {
        fail("B must resume once with TimedOut");
    }
    write_line("TURKOS/IPC-WAIT-TIMEOUT-B-WAKE: CAUSE=TIMEOUT RETRY=TIMED-OUT");
    if unsafe { (&mut *scheduler_ptr()).exit_current(0x0b) } != Ok(Some(a)) {
        fail("B exit should select A");
    }
    switch(1, 0);
    fail("B returned after exit")
}

pub(super) fn smoke_test() -> ! {
    let a = unsafe { (&mut *scheduler_ptr()).spawn() }.unwrap_or_else(|_| fail("A spawn"));
    let b = unsafe { (&mut *scheduler_ptr()).spawn() }.unwrap_or_else(|_| fail("B spawn"));
    let channel = unsafe { (&mut *core::ptr::addr_of_mut!(CHANNELS)).create(a, b) }
        .unwrap_or_else(|_| fail("channel create"));
    let a_start = unsafe { core::ptr::addr_of_mut!(A_STACK.0) as usize };
    let b_start = unsafe { core::ptr::addr_of_mut!(B_STACK.0) as usize };
    unsafe {
        write_volatile(a_start as *mut u64, CANARY_A);
        write_volatile(b_start as *mut u64, CANARY_B);
        write_volatile(core::ptr::addr_of_mut!(IDS), [a, b]);
        write_volatile(core::ptr::addr_of_mut!(CHANNEL_ID), channel);
        core::ptr::addr_of_mut!(A_CONTEXT).write(
            task::CpuContext::initial(
                (a_start + STACK_BYTES) as u64,
                task_a as *const () as usize as u64,
            )
            .unwrap_or_else(|_| fail("A context")),
        );
        core::ptr::addr_of_mut!(B_CONTEXT).write(
            task::CpuContext::initial(
                (b_start + STACK_BYTES) as u64,
                task_b as *const () as usize as u64,
            )
            .unwrap_or_else(|_| fail("B context")),
        );
    }
    if unsafe { (&mut *scheduler_ptr()).run_next() } != Ok(Some(a))
        || unsafe { (&*scheduler_ptr()).ready_len() } != 1
    {
        fail("initial schedule");
    }
    write_str("TURKOS/IPC-WAIT-START: CHANNEL=0x");
    write_hex(channel as usize);
    write_str(" A=0x");
    write_hex(a.slot as usize);
    write_str(" B=0x");
    write_hex(b.slot as usize);
    write_str("\r\n");
    unsafe { __turkos_swap_task_context(core::ptr::addr_of_mut!(BOOT_CONTEXT), context(0)) };
    if unsafe { (&*scheduler_ptr()).running() }.is_some()
        || unsafe { (&mut *scheduler_ptr()).run_next() } != Ok(None)
        || unsafe { (&*scheduler_ptr()).state(a) } != Ok(State::Exited { code: 0x0a })
        || unsafe { (&*scheduler_ptr()).state(b) } != Ok(State::Exited { code: 0x0b })
        || unsafe { read_volatile(a_start as *const u64) } != CANARY_A
        || unsafe { read_volatile(b_start as *const u64) } != CANARY_B
        || unsafe { read_volatile(core::ptr::addr_of!(A_STEPS)) } != 1
    {
        fail("final task states or stack guards");
    }
    write_line("TURKOS/IPC-WAIT-EXIT: A=0x000000000000000a B=0x000000000000000b GUARDS=PASS");
    write_line("TURKOS/IPC-WAIT: PASS");
    park()
}