//! K316/K317/K413: exact child TaskId authorization and single-use exit delivery
//! through one blocked EL0 WAIT_CHILD syscall. One CPU, masked IRQ.
use super::*;
#[cfg(feature = "serial-wait-test")]
use turkos_kernel::serial_console::{ConsoleEvent, SerialConsole};
use turkos_kernel::syscall::{self, Call};
use turkos_kernel::task::{ChildWait, Error as TaskError, State, TaskId, WakeCause};

static mut CALLS: usize = 0;

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

#[cfg(feature = "serial-wait-test")]
fn wait_for_display(parent: TaskId, child: TaskId, code: i32) {
    const RXFE: u32 = 1 << 4;
    const DATA_ERROR: u32 = 0x0f00;
    let base = UART_BASE.load(Ordering::Relaxed);
    let frequency: u64;
    unsafe {
        asm!("mrs {value}, CNTFRQ_EL0", value = out(reg) frequency, options(nomem, nostack, preserves_flags))
    };
    if frequency == 0 {
        fail("wait counter frequency");
    }
    let mut console = SerialConsole::new();
    write_line("TURKOS/WAIT: READY");
    loop {
        let start: u64;
        unsafe {
            asm!("mrs {value}, CNTPCT_EL0", value = out(reg) start, options(nomem, nostack, preserves_flags))
        };
        let deadline = start
            .checked_add(frequency.saturating_mul(5))
            .unwrap_or(u64::MAX);
        while unsafe { read_volatile((base + 0x18) as *const u32) } & RXFE != 0 {
            let now: u64;
            unsafe {
                asm!("mrs {value}, CNTPCT_EL0", value = out(reg) now, options(nomem, nostack, preserves_flags))
            };
            if now >= deadline {
                write_line("TURKOS/WAIT-ERROR: input timeout");
                park();
            }
            core::hint::spin_loop();
        }
        let data = unsafe { read_volatile(base as *const u32) };
        if data & DATA_ERROR != 0 {
            write_line("TURKOS/WAIT-ERROR: PL011 data error");
            park();
        }
        match console.feed(data as u8) {
            ConsoleEvent::Input | ConsoleEvent::SuppressedLf => {}
            ConsoleEvent::Wait { prompt } => {
                write_str("TURKOS/WAIT-ACCEPT: COMMAND=wait PROMPT=0x");
                write_hex(prompt as usize);
                write_str("\r\nTURKOS/WAIT-RESULT: PARENT=0x");
                write_hex(parent.slot as usize);
                write_str("/0x");
                write_hex(parent.generation as usize);
                write_str(" CHILD=0x");
                write_hex(child.slot as usize);
                write_str("/0x");
                write_hex(child.generation as usize);
                write_str(" CODE=0x");
                write_hex(code as usize);
                write_str(" REAPED=0x1\r\n");
                return;
            }
            _ => {
                write_line("TURKOS/WAIT-ERROR: unhandled command");
                park();
            }
        }
    }
}

fn token(index: usize) -> u64 {
    if index == 0 {
        user_task::TOKEN_A
    } else {
        user_task::TOKEN_B
    }
}

fn expected(index: usize) -> (usize, u64, u64, u64) {
    match index {
        0 => (0, syscall::WAIT_CHILD, 0, 1),
        1 => (0, syscall::WAIT_CHILD, 1, 3),
        2 => (0, syscall::WAIT_CHILD, 1, 2),
        3 => (1, syscall::WAIT_CHILD, 0, 1),
        4 => (1, syscall::EXIT, 0x2a, 0),
        5 => (0, syscall::WAIT_CHILD, 1, 2),
        6 => (0, syscall::EXIT, user_task::TOKEN_A, 0),
        _ => fail("extra syscall"),
    }
}

fn write_call(
    index: usize,
    current: usize,
    number: u64,
    slot: u64,
    generation: u64,
    result: &str,
    root: u64,
    pc: u64,
) {
    write_str("TURKOS/USER-WAIT-CALL: INDEX=0x");
    write_hex(index);
    write_str(" ID=");
    write_str(if current == 0 { "A" } else { "B" });
    write_str(" NUMBER=0x");
    write_hex(number as usize);
    write_str(" SLOT=0x");
    write_hex(slot as usize);
    write_str(" GEN=0x");
    write_hex(generation as usize);
    write_str(" RESULT=");
    write_str(result);
    write_str(" ROOT=0x");
    write_hex(root as usize);
    write_str(" ELR=0x");
    write_hex(pc as usize);
    write_str("\r\n");
}

fn write_return(index: usize, value: i64, pc: u64) {
    write_str("TURKOS/USER-WAIT-RET: INDEX=0x");
    write_hex(index);
    write_str(" X0=0x");
    write_hex(value as usize);
    write_str(" NEXT=0x");
    write_hex(pc as usize);
    write_str("\r\n");
}

fn write_switch(index: usize, from: &str, to: &str, ready: usize) {
    write_str("TURKOS/USER-WAIT-SWITCH: INDEX=0x");
    write_hex(index);
    write_str(" FROM=");
    write_str(from);
    write_str(" TO=");
    write_str(to);
    write_str(" READY=0x");
    write_hex(ready);
    write_str("\r\n");
}

#[unsafe(no_mangle)]
pub extern "C" fn kernel_return_probe(frame: *mut irq::IrqContext) -> usize {
    if frame.is_null() {
        fail("null context");
    }
    let frame = unsafe { &mut *frame };
    let esr: u64;
    let root: u64;
    unsafe {
        asm!("mrs {value}, ESR_EL1", value = out(reg) esr, options(nomem, nostack, preserves_flags));
        asm!("mrs {value}, TTBR0_EL1", value = out(reg) root, options(nomem, nostack, preserves_flags));
    }
    let count = unsafe { read_volatile(core::ptr::addr_of!(CALLS)) };
    let (current, number, arg_slot, arg_generation) = expected(count);
    let ids = unsafe { read_volatile(core::ptr::addr_of!(user_task::IDS)) };
    let roots = unsafe { read_volatile(core::ptr::addr_of!(user_task::ROOTS)) };
    let physical = unsafe { read_volatile(core::ptr::addr_of!(user_task::FRAMES)) };
    let owner = ids[current].unwrap_or_else(|| fail("task ID missing"));
    let parent = ids[0].unwrap_or_else(|| fail("parent ID missing"));
    let child = ids[1].unwrap_or_else(|| fail("child ID missing"));
    let scheduler = unsafe { &mut *core::ptr::addr_of_mut!(user_task::SCHEDULER) };
    if parent.slot != 0
        || parent.generation != 1
        || child.slot != 1
        || child.generation != 2
        || unsafe { read_volatile(core::ptr::addr_of!(__irq_stack_guard)) } != user_task::GUARD
        || scheduler.running() != Some(owner)
        || esr != 0x5600_0042
        || frame.vector_slot != 8
        || frame.spsr & 0xf != 0
        || frame.sp_el0 != mmu::USER_STACK_TOP - 16
        || frame.elr < mmu::USER_CODE_VA
        || frame.elr >= mmu::USER_CODE_VA + memory::PAGE_SIZE
        || frame.elr & 3 != 0
        || root != roots[current + 1]
        || frame.x[19] != token(current)
        || frame.x[20] != mmu::USER_DATA_VA
        || frame.x[21] != 0x55
        || user_task::value(physical[current][0]) != token(current)
        || user_task::value(physical[current][1] + memory::PAGE_SIZE - 16) != token(current)
    {
        fail("EL0 frame, owner or private page");
    }
    let call = Call {
        version: frame.x[7],
        number: frame.x[8],
        args: frame.x[..6].try_into().unwrap(),
    };
    if call.version != syscall::VERSION
        || call.number != number
        || call.args[0] != arg_slot
        || (number == syscall::WAIT_CHILD && call.args[1] != arg_generation)
    {
        fail("call order or child identity");
    }
    if number == syscall::WAIT_CHILD {
        let target = TaskId {
            slot: call.args[0] as u8,
            generation: call.args[1],
        };
        if call.args[0] > u8::MAX as u64 {
            fail("truncated child slot");
        }
        match scheduler.wait_child(owner, target) {
            Ok(ChildWait::Blocked(Some(next))) if count == 2 && next == child => {
                write_call(
                    count,
                    current,
                    number,
                    call.args[0],
                    call.args[1],
                    "BLOCKED",
                    root,
                    frame.elr,
                );
                write_str("TURKOS/USER-WAIT-BLOCK: INDEX=0x");
                write_hex(count);
                write_str(" PARENT=A CHILD=0x");
                write_hex(child.slot as usize);
                write_str("/0x");
                write_hex(child.generation as usize);
                write_str(" READY=0x");
                write_hex(scheduler.ready_len());
                write_str("\r\n");
                if scheduler.state(parent) != Ok(State::WaitingChild { child })
                    || scheduler.ready_len() != 0
                {
                    fail("parent was not blocked for exact child");
                }
                unsafe {
                    core::ptr::write(
                        core::ptr::addr_of_mut!(user_task::SAVED)
                            .cast::<irq::IrqContext>()
                            .add(0),
                        *frame,
                    )
                };
                unsafe { write_volatile(core::ptr::addr_of_mut!(CALLS), count + 1) };
                write_switch(count, "A", "B", 0);
                let mut fresh = irq::IrqContext::EMPTY;
                fresh.x[0] = user_task::TOKEN_B;
                fresh.interrupted_sp = frame.interrupted_sp;
                fresh.sp_el0 = mmu::USER_STACK_TOP;
                fresh.elr = mmu::USER_CODE_VA;
                fresh.spsr = user_task::EL0T_MASKED;
                fresh.vector_slot = 8;
                *frame = fresh;
                user_task::switch_root(roots[2]);
                return 1;
            }
            Err(TaskError::NotChild) if matches!(count, 0 | 1 | 3 | 5) => {
                let reason = match count {
                    0 => "SELF",
                    1 => "STALE",
                    3 => "FOREIGN",
                    _ => "CONSUMED",
                };
                write_str("TURKOS/USER-WAIT-REJECT: INDEX=0x");
                write_hex(count);
                write_str(" REASON=");
                write_str(reason);
                write_str(" RESULT=ECHILD\r\n");
                write_call(
                    count,
                    current,
                    number,
                    call.args[0],
                    call.args[1],
                    "ECHILD",
                    root,
                    frame.elr,
                );
                frame.x[0] = syscall::ECHILD as u64;
                write_return(count, syscall::ECHILD, frame.elr);
                unsafe { write_volatile(core::ptr::addr_of_mut!(CALLS), count + 1) };
                return 1;
            }
            _ => fail("wait ownership or transition"),
        }
    }
    if count == 4 {
        write_call(
            count,
            current,
            number,
            call.args[0],
            0,
            "EXIT",
            root,
            frame.elr,
        );
        unsafe { write_volatile(core::ptr::addr_of_mut!(CALLS), count + 1) };
        if scheduler.exit_current(0x2a) != Ok(Some(parent))
            || scheduler.state(child) != Ok(State::Exited { code: 0x2a })
            || scheduler.take_wake_cause(parent) != Ok(Some(WakeCause::ChildExited))
        {
            fail("child exit/wake/single collection");
        }
        let code = match scheduler.wait_child(parent, child) {
            Ok(ChildWait::Exited(value)) if value == 0x2a => value,
            _ => fail("child exit code collection"),
        };
        if scheduler.state(child) != Err(TaskError::InvalidTask) || scheduler.ready_len() != 0 {
            fail("child reap or ready queue");
        }
        #[cfg(feature = "serial-wait-test")]
        wait_for_display(parent, child, code);
        #[cfg(not(feature = "serial-wait-test"))]
        let _ = code;
        let saved = unsafe {
            &mut *core::ptr::addr_of_mut!(user_task::SAVED)
                .cast::<irq::IrqContext>()
                .add(0)
        };
        saved.x[0] = code as u64;
        write_line("TURKOS/USER-WAIT-WAKE: CHILD=B PARENT=A CAUSE=CHILD-EXITED CODE=0x2a");
        write_str("TURKOS/USER-WAIT-DELIVER: WAIT-INDEX=0x");
        write_hex(2);
        write_str(" X0=0x");
        write_hex(saved.x[0] as usize);
        write_str(" NEXT=0x");
        write_hex(saved.elr as usize);
        write_str(" REAPED=0x1\r\n");
        write_switch(count, "B", "A", 0);
        *frame = *saved;
        user_task::switch_root(roots[1]);
        return 1;
    }
    if count != 6 || call.args[0] != user_task::TOKEN_A {
        fail("unexpected parent exit");
    }
    write_call(
        count,
        current,
        number,
        call.args[0],
        0,
        "EXIT",
        root,
        frame.elr,
    );
    unsafe { write_volatile(core::ptr::addr_of_mut!(CALLS), count + 1) };
    if scheduler.exit_current(user_task::TOKEN_A as i32) != Ok(None)
        || scheduler.running().is_some()
        || scheduler.run_next() != Ok(None)
        || scheduler.state(parent)
            != Ok(State::Exited {
                code: user_task::TOKEN_A as i32,
            })
        || scheduler.state(child) != Err(TaskError::InvalidTask)
        || scheduler.ready_len() != 0
        || user_task::value(physical[0][0]) != user_task::TOKEN_A
        || user_task::value(physical[1][0]) != user_task::TOKEN_B
        || user_task::value(physical[0][1] + memory::PAGE_SIZE - 16) != user_task::TOKEN_A
        || user_task::value(physical[1][1] + memory::PAGE_SIZE - 16) != user_task::TOKEN_B
    {
        fail("final state/owner");
    }
    write_switch(count, "A", "NONE", 0);
    user_task::switch_root(roots[0]);
    let allocator_ptr = unsafe { read_volatile(core::ptr::addr_of!(user_task::ALLOCATOR)) };
    if allocator_ptr.is_null() {
        fail("allocator missing");
    }
    let allocator = unsafe { &mut *allocator_ptr };
    let before = allocator.remaining_pages();
    for index in 0..2 {
        let space = unsafe {
            core::ptr::replace(
                core::ptr::addr_of_mut!(user_task::SPACES)
                    .cast::<Option<mmu::UserSpace>>()
                    .add(index),
                None,
            )
        }
        .unwrap_or_else(|| fail("owner space missing"));
        unsafe { space.release_after_tlbi(allocator, roots[0]) }
            .unwrap_or_else(|_| fail("owner frame release"));
    }
    let after = allocator.remaining_pages();
    if after != before + 12
        || after != unsafe { read_volatile(core::ptr::addr_of!(user_task::INITIAL_FREE)) }
        || unsafe { read_volatile((roots[0] as *const u64).add(2)) } != 0
    {
        fail("frame reclaim accounting");
    }
    write_str("TURKOS/USER-WAIT-RECLAIM: KERNEL=0x");
    write_hex(roots[0] as usize);
    write_str(" RELEASED=0x");
    write_hex((after - before) as usize);
    write_str(" READY=0x");
    write_hex(scheduler.ready_len());
    write_str("\r\n");
    write_line("TURKOS/USER-WAIT: PASS");
    0
}