//! Live cooperative EL0 A/B context switching over the bounded Scheduler.
//! One CPU, interrupts masked and ASID 0; not timer preemption.
use super::*;
#[cfg(feature = "user-task-test")]
use turkos_kernel::syscall::{self, Call, Outcome};
#[cfg(feature = "user-task-test")]
use turkos_kernel::task::State;
use turkos_kernel::task::{Scheduler, TaskId};

pub(super) const TOKEN_A: u64 = 0xa11a;
pub(super) const TOKEN_B: u64 = 0xb22b;
pub(super) const GUARD: u64 = 0x5455_524b_4952_5147;
pub(super) const EL0T_MASKED: u64 = 0x3c0;
pub(super) static mut SCHEDULER: Scheduler = Scheduler::new();
pub(super) static mut IDS: [Option<TaskId>; 2] = [None; 2];
pub(super) static mut ROOTS: [u64; 3] = [0; 3]; // kernel, A, B
pub(super) static mut FRAMES: [[u64; 2]; 2] = [[0; 2]; 2]; // data, stack
pub(super) static mut SAVED: [irq::IrqContext; 2] = [irq::IrqContext::EMPTY; 2];
pub(super) static mut SPACES: [Option<mmu::UserSpace>; 2] = [None, None];
pub(super) static mut ALLOCATOR: *mut memory::FrameAllocator = core::ptr::null_mut();
pub(super) static mut INITIAL_FREE: u64 = 0;
#[cfg(feature = "user-task-test")]
static mut CALLS: usize = 0;

unsafe extern "C" {
    static __user_task_blob_start: u8;
    static __user_task_blob_end: u8;
}

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

#[cfg(feature = "user-task-test")]
fn token(index: usize) -> u64 {
    if index == 0 { TOKEN_A } else { TOKEN_B }
}

pub(super) fn switch_root(root: u64) {
    unsafe {
        asm!(
            "dsb ishst", "msr TTBR0_EL1, {root}", "isb",
            "tlbi vmalle1is", "dsb ish", "isb",
            root = in(reg) root,
            options(nostack)
        )
    };
    let actual: u64;
    unsafe {
        asm!("mrs {value}, TTBR0_EL1", value = out(reg) actual, options(nomem, nostack, preserves_flags))
    };
    if actual != root {
        fail("TTBR0 readback");
    }
}

pub(super) fn value(address: u64) -> u64 {
    unsafe { read_volatile(address as *const u64) }
}

pub(super) fn smoke_test(
    allocator: &mut memory::FrameAllocator,
    tables: &mmu::Tables,
    layout: mmu::Layout,
) -> ! {
    let before = allocator.remaining_pages();
    let mut actual: u64;
    unsafe {
        asm!("mrs {value}, TTBR0_EL1", value = out(reg) actual, options(nomem, nostack, preserves_flags))
    };
    if actual != tables.root {
        fail("kernel root before task launch");
    }
    let a = unsafe { tables.build_user_space(allocator, layout.ram_base, layout.ram_size) }
        .unwrap_or_else(|_| fail("A allocation"));
    let b = unsafe { tables.build_user_space(allocator, layout.ram_base, layout.ram_size) }
        .unwrap_or_else(|_| fail("B allocation"));
    if !a.contract_valid(tables.root)
        || !b.contract_valid(tables.root)
        || a.root == b.root
        || allocator.remaining_pages() + 12 != before
        || [a.code_frame(), a.data_frame(), a.stack_frame()]
            .iter()
            .any(|frame| [b.code_frame(), b.data_frame(), b.stack_frame()].contains(frame))
    {
        fail("private task roots or frames");
    }
    let start = core::ptr::addr_of!(__user_task_blob_start) as usize;
    let end = core::ptr::addr_of!(__user_task_blob_end) as usize;
    let length = end.checked_sub(start).unwrap_or(0);
    if length == 0 || length > memory::PAGE_SIZE as usize || length & 3 != 0 {
        fail("task blob length");
    }
    let ctr: usize;
    unsafe {
        asm!("mrs {value}, CTR_EL0", value = out(reg) ctr, options(nomem, nostack, preserves_flags))
    };
    let dline = 4usize << ((ctr >> 16) & 15);
    for code in [a.code_frame(), b.code_frame()] {
        unsafe { core::ptr::copy_nonoverlapping(start as *const u8, code as *mut u8, length) };
        for offset in (0..length).step_by(dline) {
            unsafe {
                asm!("dc cvau, {addr}", addr = in(reg) (code as usize + offset), options(nostack))
            };
        }
    }
    unsafe { asm!("dsb ish", "ic iallu", "dsb ish", "isb", options(nostack)) };
    let scheduler = unsafe { &mut *core::ptr::addr_of_mut!(SCHEDULER) };
    let a_id = scheduler.spawn().unwrap_or_else(|_| fail("A spawn"));
    let b_id = if cfg!(feature = "user-wait-test") {
        if scheduler.run_next() != Ok(Some(a_id)) {
            fail("parent start before child spawn");
        }
        scheduler
            .spawn_child(a_id)
            .unwrap_or_else(|_| fail("B child spawn"))
    } else {
        let id = scheduler.spawn().unwrap_or_else(|_| fail("B spawn"));
        if scheduler.run_next() != Ok(Some(a_id)) {
            fail("initial task selection");
        }
        id
    };
    #[cfg(feature = "user-handle-transfer-test")]
    user_handle::init_channel(a_id, b_id);
    #[cfg(feature = "user-handle-death-test")]
    user_handle_death::init_channel(a_id, b_id);
    if scheduler.running() != Some(a_id) || scheduler.ready_len() != 1 {
        fail("initial ready queue");
    }
    unsafe {
        write_volatile(core::ptr::addr_of_mut!(IDS), [Some(a_id), Some(b_id)]);
        write_volatile(
            core::ptr::addr_of_mut!(ROOTS),
            [tables.root, a.root, b.root],
        );
        write_volatile(
            core::ptr::addr_of_mut!(FRAMES),
            [
                [a.data_frame(), a.stack_frame()],
                [b.data_frame(), b.stack_frame()],
            ],
        );
        core::ptr::write(core::ptr::addr_of_mut!(SPACES), [Some(a), Some(b)]);
        write_volatile(core::ptr::addr_of_mut!(ALLOCATOR), allocator as *mut _);
        write_volatile(core::ptr::addr_of_mut!(INITIAL_FREE), before);
        write_volatile(core::ptr::addr_of_mut!(__irq_stack_guard), GUARD);
    }
    write_str("TURKOS/USER-TASK-ROOTS: KERNEL=0x");
    write_hex(tables.root as usize);
    write_str(" A=0x");
    write_hex(unsafe { read_volatile(core::ptr::addr_of!(ROOTS))[1] } as usize);
    write_str(" B=0x");
    write_hex(unsafe { read_volatile(core::ptr::addr_of!(ROOTS))[2] } as usize);
    write_str(" BEFORE=0x");
    write_hex(before as usize);
    write_str(" AFTER=0x");
    write_hex(allocator.remaining_pages() as usize);
    write_str("\r\nTURKOS/USER-TASK-IDS: A=0x");
    write_hex(a_id.slot as usize);
    write_str(" B=0x");
    write_hex(b_id.slot as usize);
    write_str(" READY=0x");
    write_hex(scheduler.ready_len());
    write_str("\r\n");
    #[cfg(feature = "user-wait-test")]
    {
        write_str("TURKOS/USER-WAIT-IDS: PARENT=0x");
        write_hex(a_id.slot as usize);
        write_str("/0x");
        write_hex(a_id.generation as usize);
        write_str(" CHILD=0x");
        write_hex(b_id.slot as usize);
        write_str("/0x");
        write_hex(b_id.generation as usize);
        write_str("\r\n");
    }
    let roots = unsafe { read_volatile(core::ptr::addr_of!(ROOTS)) };
    switch_root(roots[1]);
    write_line(if cfg!(feature = "user-wait-test") {
        "TURKOS/USER-WAIT-ENTRY: A EL0"
    } else if cfg!(feature = "user-handle-death-test") {
        "TURKOS/USER-HANDLE-DEATH-ENTRY: A EL0"
    } else if cfg!(any(
        feature = "user-handle-test",
        feature = "user-handle-transfer-test"
    )) {
        "TURKOS/USER-HANDLE-ENTRY: A EL0"
    } else {
        "TURKOS/USER-TASK-ENTRY: A EL0"
    });
    unsafe {
        asm!(
            "msr SP_EL0, {stack}", "msr ELR_EL1, {entry}",
            "msr SPSR_EL1, {mode}", "isb", "eret",
            stack = in(reg) mmu::USER_STACK_TOP,
            entry = in(reg) mmu::USER_CODE_VA,
            mode = in(reg) EL0T_MASKED,
            in("x0") TOKEN_A,
            options(noreturn)
        )
    }
}

#[cfg(feature = "user-task-test")]
#[unsafe(no_mangle)]
pub extern "C" fn kernel_return_probe(frame: *mut irq::IrqContext) -> usize {
    if frame.is_null() {
        fail("null task frame");
    }
    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 roots = unsafe { read_volatile(core::ptr::addr_of!(ROOTS)) };
    let ids = unsafe { read_volatile(core::ptr::addr_of!(IDS)) };
    let physical = unsafe { read_volatile(core::ptr::addr_of!(FRAMES)) };
    let current = [0usize, 1, 0, 1]
        .get(count)
        .copied()
        .unwrap_or_else(|| fail("extra EL0 call"));
    let id = ids[current].unwrap_or_else(|| fail("missing task ID"));
    let scheduler = unsafe { &mut *core::ptr::addr_of_mut!(SCHEDULER) };
    if unsafe { read_volatile(core::ptr::addr_of!(__irq_stack_guard)) } != GUARD
        || frame.vector_slot != 8
        || esr != 0x5600_0042
        || 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
        || root != roots[current + 1]
        || scheduler.running() != Some(id)
        || frame.x[19] != token(current)
        || frame.x[20] != mmu::USER_DATA_VA
        || frame.x[21] != 0x55
        || value(physical[current][0]) != token(current)
        || value(physical[current][1] + memory::PAGE_SIZE - 16) != token(current)
    {
        fail("EL0 frame, root, register, or private page");
    }
    let call = Call {
        version: frame.x[7],
        number: frame.x[8],
        args: frame.x[..6].try_into().unwrap(),
    };
    let expected_number = if count < 2 {
        syscall::YIELD
    } else {
        syscall::EXIT
    };
    let expected_outcome = if count < 2 {
        Outcome::Yield
    } else {
        Outcome::Exit(token(current))
    };
    if call.number != expected_number || syscall::dispatch(call) != expected_outcome {
        fail("task syscall ABI or order");
    }
    if count == 2 && scheduler.state(ids[1].unwrap()) != Ok(State::Ready) {
        fail("peer was not ready before A exit");
    }
    write_str("TURKOS/USER-TASK-CALL: INDEX=0x");
    write_hex(count);
    write_str(" ID=");
    write_str(if current == 0 { "A" } else { "B" });
    write_str(" NUMBER=0x");
    write_hex(call.number as usize);
    write_str(" ROOT=0x");
    write_hex(root as usize);
    write_str(" ESR=0x");
    write_hex(esr as usize);
    write_str(" SLOT=0x");
    write_hex(frame.vector_slot as usize);
    write_str(" ELR=0x");
    write_hex(frame.elr as usize);
    write_str(" SP=0x");
    write_hex(frame.sp_el0 as usize);
    write_str(" TOKEN=0x");
    write_hex(token(current) as usize);
    write_str("\r\n");
    let next = if count < 2 {
        frame.x[0] = 0;
        unsafe {
            core::ptr::write(
                core::ptr::addr_of_mut!(SAVED)
                    .cast::<irq::IrqContext>()
                    .add(current),
                *frame,
            )
        };
        scheduler
            .yield_current(false)
            .unwrap_or_else(|_| fail("yield transition"))
    } else {
        scheduler
            .exit_current(token(current) as i32)
            .unwrap_or_else(|_| fail("exit transition"))
    };
    let expected_next = [
        Some(ids[1].unwrap()),
        Some(ids[0].unwrap()),
        Some(ids[1].unwrap()),
        None,
    ][count];
    if next != expected_next || scheduler.ready_len() != if count < 2 { 1 } else { 0 } {
        fail("ready queue or next task");
    }
    write_str("TURKOS/USER-TASK-SWITCH: INDEX=0x");
    write_hex(count);
    write_str(" FROM=");
    write_str(if current == 0 { "A" } else { "B" });
    write_str(" TO=");
    write_str(match count {
        0 | 2 => "B",
        1 => "A",
        _ => "NONE",
    });
    write_str(" READY=0x");
    write_hex(scheduler.ready_len());
    write_str(" A-STATE=0x");
    write_hex(match scheduler.state(ids[0].unwrap()).unwrap() {
        State::Exited { .. } => 3,
        State::Ready => 1,
        State::Running => 2,
        _ => 0,
    });
    write_str(" B-STATE=0x");
    write_hex(match scheduler.state(ids[1].unwrap()).unwrap() {
        State::Exited { .. } => 3,
        State::Ready => 1,
        State::Running => 2,
        _ => 0,
    });
    write_str("\r\n");
    unsafe { write_volatile(core::ptr::addr_of_mut!(CALLS), count + 1) };
    if let Some(next_id) = next {
        let next_index = if next_id == ids[0].unwrap() { 0 } else { 1 };
        if count == 0 {
            let mut fresh = irq::IrqContext::EMPTY;
            fresh.x[0] = TOKEN_B;
            fresh.interrupted_sp = frame.interrupted_sp;
            fresh.sp_el0 = mmu::USER_STACK_TOP;
            fresh.elr = mmu::USER_CODE_VA;
            fresh.spsr = EL0T_MASKED;
            fresh.vector_slot = 8;
            *frame = fresh;
        } else {
            *frame = unsafe {
                core::ptr::read(
                    core::ptr::addr_of!(SAVED)
                        .cast::<irq::IrqContext>()
                        .add(next_index),
                )
            };
        }
        switch_root(roots[next_index + 1]);
        return 1;
    }
    if scheduler.running().is_some()
        || scheduler.run_next() != Ok(None)
        || scheduler.state(ids[0].unwrap())
            != Ok(State::Exited {
                code: TOKEN_A as i32,
            })
        || scheduler.state(ids[1].unwrap())
            != Ok(State::Exited {
                code: TOKEN_B as i32,
            })
        || value(physical[0][0]) != TOKEN_A
        || value(physical[1][0]) != TOKEN_B
        || value(physical[0][1] + memory::PAGE_SIZE - 16) != TOKEN_A
        || value(physical[1][1] + memory::PAGE_SIZE - 16) != TOKEN_B
    {
        fail("exited task selected or private memory changed");
    }
    switch_root(roots[0]);
    let allocator_ptr = unsafe { read_volatile(core::ptr::addr_of!(ALLOCATOR)) };
    if allocator_ptr.is_null() {
        fail("allocator missing");
    }
    let allocator = unsafe { &mut *allocator_ptr };
    let before_reclaim = allocator.remaining_pages();
    for index in 0..2 {
        let space = unsafe {
            core::ptr::replace(
                core::ptr::addr_of_mut!(SPACES)
                    .cast::<Option<mmu::UserSpace>>()
                    .add(index),
                None,
            )
        }
        .unwrap_or_else(|| fail("task owner missing"));
        unsafe { space.release_after_tlbi(allocator, roots[0]) }
            .unwrap_or_else(|_| fail("task frame reclaim"));
    }
    let after_reclaim = allocator.remaining_pages();
    if before_reclaim + 12 != after_reclaim
        || after_reclaim != unsafe { read_volatile(core::ptr::addr_of!(INITIAL_FREE)) }
        || unsafe { read_volatile((roots[0] as *const u64).add(2)) } != 0
    {
        fail("reclaim accounting or kernel user slot");
    }
    write_str("TURKOS/USER-TASK-RECLAIM: KERNEL=0x");
    write_hex(roots[0] as usize);
    write_str(" RELEASED=0x");
    write_hex((after_reclaim - before_reclaim) as usize);
    write_str(" FREE=0x");
    write_hex(after_reclaim as usize);
    write_str("\r\n");
    write_line("TURKOS/USER-TASK: PASS");
    0
}