//! Two sequential ELF-backed EL0 programs. A is reclaimed before B is
//! installed; a malformed B must leave its owned pages untouched.
use super::*;
#[cfg(feature = "serial-run-test")]
use turkos_kernel::serial_console::{ConsoleEvent, SerialConsole};
use turkos_kernel::{elf, program_hash, syscall, user_copy};

const GUARD: u64 = 0x454c_4653_4551_4744;
const PROGRAM_LEN: usize = 56;
static mut ROOTS: [u64; 3] = [0; 3];
static mut FRAMES: [[u64; 3]; 2] = [[0; 3]; 2];
static mut SPACES: [Option<mmu::UserSpace>; 2] = [None, None];
static mut ALLOCATOR: *mut memory::FrameAllocator = core::ptr::null_mut();
static mut INITIAL_FREE: u64 = 0;
static mut INITIAL_SP: [u64; 2] = [0; 2];
static mut CALLS: usize = 0;

unsafe extern "C" {
    static __elf_hello_blob_start: u8;
    static __elf_hello_blob_end: u8;
}

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

fn fixture(payload: &[u8; 5]) -> [u8; 512] {
    let mut file = [0u8; 512];
    file[..20].copy_from_slice(&[
        0x7f, b'E', b'L', b'F', 2, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0xb7, 0,
    ]);
    file[20..24].copy_from_slice(&1u32.to_le_bytes());
    file[24..32].copy_from_slice(&mmu::USER_CODE_VA.to_le_bytes());
    file[32..40].copy_from_slice(&64u64.to_le_bytes());
    file[52..54].copy_from_slice(&64u16.to_le_bytes());
    file[54..56].copy_from_slice(&56u16.to_le_bytes());
    file[56..58].copy_from_slice(&2u16.to_le_bytes());
    file[64..68].copy_from_slice(&1u32.to_le_bytes());
    file[68..72].copy_from_slice(&5u32.to_le_bytes());
    file[72..80].copy_from_slice(&320u64.to_le_bytes());
    file[80..88].copy_from_slice(&mmu::USER_CODE_VA.to_le_bytes());
    file[96..104].copy_from_slice(&56u64.to_le_bytes());
    file[104..112].copy_from_slice(&64u64.to_le_bytes());
    file[112..120].copy_from_slice(&1u64.to_le_bytes());
    file[120..124].copy_from_slice(&1u32.to_le_bytes());
    file[124..128].copy_from_slice(&6u32.to_le_bytes());
    file[128..136].copy_from_slice(&448u64.to_le_bytes());
    file[136..144].copy_from_slice(&mmu::USER_DATA_VA.to_le_bytes());
    file[152..160].copy_from_slice(&5u64.to_le_bytes());
    file[160..168].copy_from_slice(&8u64.to_le_bytes());
    file[168..176].copy_from_slice(&1u64.to_le_bytes());
    file[448..453].copy_from_slice(payload);
    let start = core::ptr::addr_of!(__elf_hello_blob_start) as usize;
    let end = core::ptr::addr_of!(__elf_hello_blob_end) as usize;
    if end.checked_sub(start) != Some(PROGRAM_LEN) {
        fail("ELF program instruction length");
    }
    unsafe {
        core::ptr::copy_nonoverlapping(
            start as *const u8,
            file[320..320 + PROGRAM_LEN].as_mut_ptr(),
            PROGRAM_LEN,
        )
    };
    file
}

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");
    }
}

fn sync_code(code_frame: u64) {
    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 offset in (0..PROGRAM_LEN).step_by(dline) {
        unsafe {
            asm!("dc cvau, {addr}", addr = in(reg) (code_frame as usize + offset), options(nostack))
        };
    }
    unsafe { asm!("dsb ish", "ic iallu", "dsb ish", "isb", options(nostack)) };
}

#[cfg(feature = "serial-run-test")]
fn wait_for_run() {
    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("run counter frequency");
    }
    let mut console = SerialConsole::run_only();
    write_line("TURKOS/RUN-PROMPT: NUMBER=0x1");
    write_line("TURKOS/RUN: 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/RUN-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/RUN-ERROR: PL011 data error");
            park();
        }
        match console.feed(data as u8) {
            ConsoleEvent::Input | ConsoleEvent::SuppressedLf => {}
            ConsoleEvent::Run { prompt } => {
                write_str("TURKOS/RUN-ACCEPT: COMMAND=run PROMPT=0x");
                write_hex(prompt as usize);
                write_str("\r\n");
                return;
            }
            _ => {
                write_line("TURKOS/RUN-ERROR: unhandled command");
                park();
            }
        }
    }
}

fn load_program(file: &[u8; 512], space: &mmu::UserSpace, kernel_root: u64, id: &str) -> u64 {
    let stack = unsafe {
        core::slice::from_raw_parts_mut(space.stack_frame() as *mut u8, memory::PAGE_SIZE as usize)
    };
    let initial = elf::prepare_user_stack(stack, &[id.as_bytes()])
        .unwrap_or_else(|_| fail("ELF sequence stack"));
    let loaded = unsafe { elf::load_into_user_space(file, space, kernel_root) }
        .unwrap_or_else(|_| fail("ELF sequence install"));
    if loaded.entry != mmu::USER_CODE_VA
        || loaded.code_bytes != PROGRAM_LEN
        || loaded.data_bytes != 5
    {
        fail("ELF sequence installed size");
    }
    let code =
        unsafe { core::slice::from_raw_parts(space.code_frame() as *const u8, loaded.code_bytes) };
    let data =
        unsafe { core::slice::from_raw_parts(space.data_frame() as *const u8, loaded.data_bytes) };
    let digest = program_hash::digest(loaded.entry, code, data);
    if digest != program_hash::digest(loaded.entry, &file[320..376], &file[448..453]) {
        fail("ELF sequence installed hash");
    }
    sync_code(space.code_frame());
    let hex = program_hash::hex_lower(&digest);
    write_str("TURKOS/ELF-SEQUENCE-START: ID=");
    write_str(id);
    write_str(" ENTRY=0x");
    write_hex(loaded.entry as usize);
    write_str(" SP=0x");
    write_hex(initial.sp as usize);
    write_str(" SHA256=");
    write_str(core::str::from_utf8(&hex).unwrap_or_else(|_| fail("hash hex")));
    write_str("\r\n");
    initial.sp
}

pub(super) fn smoke_test(
    allocator: &mut memory::FrameAllocator,
    tables: &mmu::Tables,
    ram: mmu::Layout,
) -> ! {
    let before = allocator.remaining_pages();
    let a = unsafe { tables.build_user_space(allocator, ram.ram_base, ram.ram_size) }
        .unwrap_or_else(|_| fail("A private root"));
    let b = unsafe { tables.build_user_space(allocator, ram.ram_base, ram.ram_size) }
        .unwrap_or_else(|_| fail("B private root"));
    let a_frames = [a.code_frame(), a.data_frame(), a.stack_frame()];
    let b_frames = [b.code_frame(), b.data_frame(), b.stack_frame()];
    if !a.contract_valid(tables.root)
        || !b.contract_valid(tables.root)
        || a.root == b.root
        || a_frames.iter().any(|frame| b_frames.contains(frame))
        || allocator.remaining_pages() + 12 != before
    {
        fail("ELF sequence private ownership");
    }
    unsafe {
        core::slice::from_raw_parts_mut(b.code_frame() as *mut u8, memory::PAGE_SIZE as usize)
            .fill(0xa5);
        core::slice::from_raw_parts_mut(b.data_frame() as *mut u8, memory::PAGE_SIZE as usize)
            .fill(0x5a);
    }
    let file = fixture(b"HELLO");
    let sp = load_program(&file, &a, tables.root, "A");
    unsafe {
        write_volatile(
            core::ptr::addr_of_mut!(ROOTS),
            [tables.root, a.root, b.root],
        );
        write_volatile(core::ptr::addr_of_mut!(FRAMES), [a_frames, b_frames]);
        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!(INITIAL_SP), [sp, 0]);
        write_volatile(core::ptr::addr_of_mut!(__irq_stack_guard), GUARD);
    }
    let roots = unsafe { read_volatile(core::ptr::addr_of!(ROOTS)) };
    write_str("TURKOS/ELF-SEQUENCE-ROOTS: KERNEL=0x");
    write_hex(roots[0] as usize);
    write_str(" A=0x");
    write_hex(roots[1] as usize);
    write_str(" B=0x");
    write_hex(roots[2] as usize);
    write_str(" BEFORE=0x");
    write_hex(before as usize);
    write_str(" ALLOCATED=0x");
    write_hex(allocator.remaining_pages() as usize);
    write_str("\r\nTURKOS/ELF-SEQUENCE-FRAMES: A_CODE=0x");
    write_hex(a_frames[0] as usize);
    write_str(" A_DATA=0x");
    write_hex(a_frames[1] as usize);
    write_str(" A_STACK=0x");
    write_hex(a_frames[2] as usize);
    write_str(" B_CODE=0x");
    write_hex(b_frames[0] as usize);
    write_str(" B_DATA=0x");
    write_hex(b_frames[1] as usize);
    write_str(" B_STACK=0x");
    write_hex(b_frames[2] as usize);
    write_str("\r\n");
    switch_root(roots[1]);
    write_line("TURKOS/ELF-SEQUENCE-EL0: A");
    unsafe {
        asm!(
            "msr SP_EL0, {stack}", "msr ELR_EL1, {entry}", "msr SPSR_EL1, {mode}", "isb", "eret",
            stack = in(reg) sp,
            entry = in(reg) mmu::USER_CODE_VA,
            mode = in(reg) 0x3c0usize,
            options(noreturn)
        )
    }
}

fn release(index: usize, kernel_root: u64) -> u64 {
    let ptr = unsafe { read_volatile(core::ptr::addr_of!(ALLOCATOR)) };
    if ptr.is_null() {
        fail("allocator pointer");
    }
    let allocator = unsafe { &mut *ptr };
    let owner = unsafe {
        core::ptr::replace(
            core::ptr::addr_of_mut!(SPACES)
                .cast::<Option<mmu::UserSpace>>()
                .add(index),
            None,
        )
    }
    .unwrap_or_else(|| fail("ELF owner missing"));
    unsafe { owner.release_after_tlbi(allocator, kernel_root) }
        .unwrap_or_else(|_| fail("ELF owner reclaim"));
    allocator.remaining_pages()
}

#[unsafe(no_mangle)]
pub extern "C" fn kernel_return_probe(frame: *mut irq::IrqContext) -> usize {
    if frame.is_null() {
        fail("null ELF sequence frame");
    }
    let frame = unsafe { &mut *frame };
    let count = unsafe { read_volatile(core::ptr::addr_of!(CALLS)) };
    if count > 3 {
        fail("extra ELF sequence syscall");
    }
    let index = count / 2;
    let id = if index == 0 { "A" } else { "B" };
    let roots = unsafe { read_volatile(core::ptr::addr_of!(ROOTS)) };
    let frames = unsafe { read_volatile(core::ptr::addr_of!(FRAMES)) };
    let stacks = unsafe { read_volatile(core::ptr::addr_of!(INITIAL_SP)) };
    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));
    }
    if esr != 0x5600_0042
        || frame.vector_slot != 8
        || frame.spsr & 15 != 0
        || frame.sp_el0 != stacks[index]
        || frame.elr != mmu::USER_CODE_VA + if count & 1 == 0 { 24 } else { 48 }
        || root != roots[index + 1]
        || unsafe { read_volatile(core::ptr::addr_of!(__irq_stack_guard)) } != GUARD
        || frame.x[7] != syscall::VERSION
    {
        fail("ELF sequence frame/root/ABI");
    }
    write_str("TURKOS/ELF-SEQUENCE-CALL: INDEX=0x");
    write_hex(count);
    write_str(" ID=");
    write_str(id);
    write_str(" NUMBER=0x");
    write_hex(frame.x[8] as usize);
    write_str(" ROOT=0x");
    write_hex(root as usize);
    write_str(" ELR=0x");
    write_hex(frame.elr as usize);
    write_str("\r\n");
    if count & 1 == 0 {
        if frame.x[8] != syscall::SERIAL_WRITE || frame.x[0] != mmu::USER_DATA_VA || frame.x[1] != 5
        {
            fail("ELF sequence write ABI");
        }
        let pages = user_copy::UserPages {
            code_frame: frames[index][0],
            data_frame: frames[index][1],
            stack_frame: frames[index][2],
        };
        let mut scratch = [0u8; 5];
        if unsafe { pages.copy_from(frame.x[0], &mut scratch) } != Ok(5) {
            fail("ELF sequence user copy");
        }
        let text = user_copy::serial_text(&scratch).unwrap_or_else(|_| fail("serial text"));
        if text != if index == 0 { "HELLO" } else { "WORLD" } {
            fail("ELF sequence payload");
        }
        write_str("TURKOS/ELF-SEQUENCE-TEXT: ID=");
        write_str(id);
        write_str(" VALUE=");
        write_line(text);
        frame.x[0] = 5;
        unsafe { write_volatile(core::ptr::addr_of_mut!(CALLS), count + 1) };
        1
    } else {
        if frame.x[8] != syscall::EXIT || frame.x[0] != 0 {
            fail("ELF sequence exit ABI");
        }
        write_str("TURKOS/ELF-SEQUENCE-EXIT: INDEX=0x");
        write_hex(count);
        write_str(" ID=");
        write_line(id);
        switch_root(roots[0]);
        let free = release(index, roots[0]);
        let initial = unsafe { read_volatile(core::ptr::addr_of!(INITIAL_FREE)) };
        if free != initial - if index == 0 { 6 } else { 0 } {
            fail("ELF sequence reclaim count");
        }
        write_str("TURKOS/ELF-SEQUENCE-RECLAIM: ID=");
        write_str(id);
        write_str(" FREE=0x");
        write_hex(free as usize);
        write_str("\r\n");
        if index == 1 {
            write_line("TURKOS/ELF-SEQUENCE: PASS");
            park()
        }
        let b_ptr = core::ptr::addr_of_mut!(SPACES).cast::<Option<mmu::UserSpace>>();
        let b = unsafe { (&*b_ptr.add(1)).as_ref() }.unwrap_or_else(|| fail("B owner missing"));
        let code = unsafe {
            core::slice::from_raw_parts(b.code_frame() as *const u8, memory::PAGE_SIZE as usize)
        };
        let data = unsafe {
            core::slice::from_raw_parts(b.data_frame() as *const u8, memory::PAGE_SIZE as usize)
        };
        if code.iter().any(|byte| *byte != 0xa5) || data.iter().any(|byte| *byte != 0x5a) {
            fail("B sentinel before bad ELF");
        }
        let mut bad = fixture(b"WORLD");
        bad[96..104].copy_from_slice(&65u64.to_le_bytes());
        if unsafe { elf::load_into_user_space(&bad, b, roots[0]) }
            != Err(elf::Error::FileLargerThanMemory)
            || code.iter().any(|byte| *byte != 0xa5)
            || data.iter().any(|byte| *byte != 0x5a)
        {
            fail("bad ELF touched B pages");
        }
        write_line("TURKOS/ELF-SEQUENCE-BAD: REJECT=FileLargerThanMemory UNCHANGED=PASS");
        #[cfg(feature = "serial-run-test")]
        wait_for_run();
        let good = fixture(b"WORLD");
        let sp = load_program(&good, b, roots[0], "B");
        unsafe {
            write_volatile(core::ptr::addr_of_mut!(INITIAL_SP), [stacks[0], sp]);
            write_volatile(core::ptr::addr_of_mut!(CALLS), 2);
        }
        frame.x = [0; 31];
        frame.q = [[0; 2]; 32];
        frame.fpcr = 0;
        frame.fpsr = 0;
        frame.sp_el0 = sp;
        frame.elr = mmu::USER_CODE_VA;
        frame.spsr = 0x3c0;
        switch_root(roots[2]);
        write_line("TURKOS/ELF-SEQUENCE-EL0: B");
        1
    }
}