//! K419: two real EL1 task contexts contend for one complete UART line.
//! The unlocked feature is a real guest counterexample with interleaved bytes.
use super::*;
use turkos_kernel::task::{CpuContext, Scheduler, State};

const STACK_BYTES: usize = 16 * 1024;
const A_GUARD: u64 = 0x4c4f_434b_5f41_3031;
const B_GUARD: u64 = 0x4c4f_434b_5f42_3031;

#[repr(align(4096))]
struct Stack([u8; STACK_BYTES]);

static mut A_STACK: Stack = Stack([0; STACK_BYTES]);
static mut B_STACK: Stack = Stack([0; STACK_BYTES]);
static mut CONTEXTS: [CpuContext; 2] = [CpuContext::EMPTY; 2];
static mut BOOT_CONTEXT: CpuContext = CpuContext::EMPTY;
static mut SCHEDULER: Scheduler = Scheduler::new();
static mut SWAPS: usize = 0;
static mut BUSY_A: usize = 0;
static mut BUSY_B: usize = 0;

fn fail(reason: &str) -> ! {
    // A failed task may still own the line. Never recursively acquire it.
    write_str("TURKOS/LOCK-FAIL: ");
    write_str(reason);
    write_str("\r\n");
    park()
}

fn context(slot: usize) -> *mut CpuContext {
    if slot >= 2 {
        fail("invalid task slot");
    }
    unsafe {
        core::ptr::addr_of_mut!(CONTEXTS)
            .cast::<CpuContext>()
            .add(slot)
    }
}

fn yield_task() {
    let (old, next) = unsafe {
        let scheduler = &mut *core::ptr::addr_of_mut!(SCHEDULER);
        let old = scheduler
            .running()
            .unwrap_or_else(|| fail("yield without task"));
        let next = scheduler
            .yield_current(false)
            .unwrap_or_else(|_| fail("yield rejected"))
            .unwrap_or_else(|| fail("no second ready task"));
        (old.slot as usize, next.slot as usize)
    };
    if old == next {
        fail("yield did not switch task");
    }
    unsafe {
        let swaps = read_volatile(core::ptr::addr_of!(SWAPS));
        write_volatile(core::ptr::addr_of_mut!(SWAPS), swaps + 1);
        __turkos_swap_task_context(context(old), context(next));
    }
}

fn exit_task(code: i32) -> ! {
    let (old, next) = unsafe {
        let scheduler = &mut *core::ptr::addr_of_mut!(SCHEDULER);
        let old = scheduler
            .running()
            .unwrap_or_else(|| fail("exit without task"));
        let next = scheduler
            .exit_current(code)
            .unwrap_or_else(|_| fail("exit rejected"));
        (old.slot as usize, next.map(|id| id.slot as usize))
    };
    unsafe {
        let swaps = read_volatile(core::ptr::addr_of!(SWAPS));
        write_volatile(core::ptr::addr_of_mut!(SWAPS), swaps + 1);
        let destination = match next {
            Some(slot) => context(slot).cast_const(),
            None => core::ptr::addr_of!(BOOT_CONTEXT),
        };
        __turkos_swap_task_context(context(old), destination);
    }
    fail("exited task resumed")
}

extern "C" fn locked_a() -> ! {
    let guard = SERIAL_LINE_LOCK
        .try_lock()
        .unwrap_or_else(|| fail("A first lock"));
    write_str("TURKOS/LOCK-LINE:A0-START-");
    yield_task();
    if unsafe { read_volatile(core::ptr::addr_of!(BUSY_B)) } != 1 {
        fail("B did not observe A ownership");
    }
    write_str("A0-END\r\n");
    drop(guard);
    yield_task();
    if let Some(unexpected) = SERIAL_LINE_LOCK.try_lock() {
        drop(unexpected);
        fail("A entered B partial line");
    }
    unsafe { write_volatile(core::ptr::addr_of_mut!(BUSY_A), 1) };
    yield_task();
    if SERIAL_LINE_LOCK.is_locked() {
        fail("B lock not released");
    }
    write_line("TURKOS/LOCK-LINE:A1-WHOLE");
    exit_task(0x0a)
}

extern "C" fn locked_b() -> ! {
    if let Some(unexpected) = SERIAL_LINE_LOCK.try_lock() {
        drop(unexpected);
        fail("B entered A partial line");
    }
    unsafe { write_volatile(core::ptr::addr_of_mut!(BUSY_B), 1) };
    yield_task();
    if SERIAL_LINE_LOCK.is_locked() {
        fail("A lock not released");
    }
    write_line("TURKOS/LOCK-LINE:B0-WHOLE");
    let guard = SERIAL_LINE_LOCK
        .try_lock()
        .unwrap_or_else(|| fail("B second lock"));
    write_str("TURKOS/LOCK-LINE:B1-START-");
    yield_task();
    if unsafe { read_volatile(core::ptr::addr_of!(BUSY_A)) } != 1 {
        fail("A did not observe B ownership");
    }
    write_str("B1-END\r\n");
    drop(guard);
    yield_task();
    write_line("TURKOS/LOCK-LINE:B2-WHOLE");
    exit_task(0x0b)
}

extern "C" fn unlocked_a() -> ! {
    write_str("TURKOS/LOCK-LINE:A0-START-");
    yield_task();
    write_str("A0-END\r\n");
    exit_task(0x0a)
}

extern "C" fn unlocked_b() -> ! {
    write_line("TURKOS/LOCK-LINE:B0-WHOLE");
    yield_task();
    exit_task(0x0b)
}

pub(super) fn smoke_test() {
    unsafe { asm!("msr daifset, #2", "isb", options(nostack, preserves_flags)) };
    let unlocked = cfg!(feature = "serial-line-unlocked-test");
    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 };
    let a_top = a_start + STACK_BYTES;
    let b_top = b_start + STACK_BYTES;
    if a_start & 4095 != 0 || b_start & 4095 != 0 || (a_start < b_top && b_start < a_top) {
        fail("task stack bounds");
    }
    unsafe {
        write_volatile(a_start as *mut u64, A_GUARD);
        write_volatile(b_start as *mut u64, B_GUARD);
    }
    let (a, b) = unsafe {
        let scheduler = &mut *core::ptr::addr_of_mut!(SCHEDULER);
        let a = scheduler.spawn().unwrap_or_else(|_| fail("spawn A"));
        let b = scheduler.spawn().unwrap_or_else(|_| fail("spawn B"));
        (a, b)
    };
    if a.slot != 0 || b.slot != 1 || a.generation == b.generation {
        fail("distinct task identities");
    }
    let a_entry = if unlocked { unlocked_a } else { locked_a };
    let b_entry = if unlocked { unlocked_b } else { locked_b };
    let a_context = CpuContext::initial(a_top as u64, a_entry as *const () as u64)
        .unwrap_or_else(|_| fail("A context"));
    let b_context = CpuContext::initial(b_top as u64, b_entry as *const () as u64)
        .unwrap_or_else(|_| fail("B context"));
    unsafe {
        context(0).write(a_context);
        context(1).write(b_context);
    }
    write_str("TURKOS/LOCK-IDS: A-SLOT=0x");
    write_hex(a.slot as usize);
    write_str(" A-GEN=0x");
    write_hex(a.generation as usize);
    write_str(" B-SLOT=0x");
    write_hex(b.slot as usize);
    write_str(" B-GEN=0x");
    write_hex(b.generation as usize);
    write_str("\r\n");
    write_line("TURKOS/LOCK: READY");
    if unsafe { (&mut *core::ptr::addr_of_mut!(SCHEDULER)).run_next() } != Ok(Some(a)) {
        fail("first task");
    }
    unsafe { __turkos_swap_task_context(core::ptr::addr_of_mut!(BOOT_CONTEXT), context(0)) };
    let scheduler = unsafe { &mut *core::ptr::addr_of_mut!(SCHEDULER) };
    let swaps = unsafe { read_volatile(core::ptr::addr_of!(SWAPS)) };
    let busy_a = unsafe { read_volatile(core::ptr::addr_of!(BUSY_A)) };
    let busy_b = unsafe { read_volatile(core::ptr::addr_of!(BUSY_B)) };
    if scheduler.running().is_some()
        || scheduler.run_next() != Ok(None)
        || scheduler.state(a) != Ok(State::Exited { code: 0x0a })
        || scheduler.state(b) != Ok(State::Exited { code: 0x0b })
        || unsafe { read_volatile(a_start as *const u64) } != A_GUARD
        || unsafe { read_volatile(b_start as *const u64) } != B_GUARD
        || SERIAL_LINE_LOCK.is_locked()
    {
        fail("task exit, guards or lock release");
    }
    if unlocked {
        if swaps != 4 || busy_a != 0 || busy_b != 0 {
            fail("unlocked control accounting");
        }
        write_line("TURKOS/LOCK-UNLOCKED: MIXED");
        park()
    }
    if swaps != 8 || busy_a != 1 || busy_b != 1 {
        fail("locked task accounting");
    }
    write_str("TURKOS/LOCK-STATE: SWAPS=0x");
    write_hex(swaps);
    write_str(" BUSY-A=0x");
    write_hex(busy_a);
    write_str(" BUSY-B=0x");
    write_hex(busy_b);
    write_str(" GUARDS=0x");
    write_hex(2);
    write_str("\r\n");
    write_line("TURKOS/LOCK: PASS LINES=0x5 TASKS=0x2 BUSY=0x2 SWAPS=0x8");
    park()
}