//! K315: a live EL0 task exit reaps its exact channel objects and all handles.
//! One CPU, masked IRQ, two private TTBR0 roots, one queued message.
use super::*;
use turkos_kernel::handle::{Resource, Table};
use turkos_kernel::ipc::{self, ChannelTable};
use turkos_kernel::syscall::{self, Call, Outcome};
use turkos_kernel::task::{State, TaskId};
use turkos_kernel::user_copy::UserPages;

static mut HANDLES: Table = Table::new();
static mut CHANNELS: ChannelTable = ChannelTable::new();
static mut CHANNEL_ID: u64 = 0;
static mut CALLS: usize = 0;

pub(super) fn init_channel(producer: TaskId, consumer: TaskId) {
    let channels = unsafe { &mut *core::ptr::addr_of_mut!(CHANNELS) };
    let id = channels
        .create(producer, consumer)
        .unwrap_or_else(|_| fail("channel create"));
    if id != 1 || channels.live_channels() != 1 {
        fail("channel ID/count");
    }
    unsafe { write_volatile(core::ptr::addr_of_mut!(CHANNEL_ID), id) };
}

fn fail(reason: &str) -> ! {
    write_str("TURKOS/USER-HANDLE-DEATH-ERROR: ");
    write_line(reason);
    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, Option<i64>, usize) {
    use syscall::{EXIT, HANDLE_OPEN, HANDLE_QUERY, IPC_SEND, YIELD};
    match index {
        0 => (0, HANDLE_OPEN, 3, Some(1), 1),
        1 => (0, IPC_SEND, 1, Some(4), 1),
        2 => (0, YIELD, 1, Some(0), 1),
        3 => (1, HANDLE_QUERY, 1, Some(syscall::EBADF), 1),
        4 => (1, HANDLE_OPEN, 2, Some(2), 2),
        5 => (1, YIELD, 2, Some(0), 2),
        6 => (0, EXIT, user_task::TOKEN_A, None, 1),
        7 => (1, HANDLE_QUERY, 1, Some(syscall::EBADF), 1),
        8 => (1, HANDLE_QUERY, 2, Some(2), 1),
        9 => (1, HANDLE_OPEN, 1, Some(3), 2),
        10 => (1, HANDLE_QUERY, 3, Some(1), 2),
        11 => (1, EXIT, user_task::TOKEN_B, None, 0),
        _ => fail("extra syscall"),
    }
}

fn page_for(index: usize) -> UserPages {
    let spaces = unsafe { &*core::ptr::addr_of!(user_task::SPACES) };
    let space = spaces[index]
        .as_ref()
        .unwrap_or_else(|| fail("private space"));
    UserPages {
        code_frame: space.code_frame(),
        data_frame: space.data_frame(),
        stack_frame: space.stack_frame(),
    }
}

fn exit_cleanup(
    owner: TaskId,
    count: usize,
    table: &mut Table,
    scheduler: &mut turkos_kernel::task::Scheduler,
) {
    let channels = unsafe { &mut *core::ptr::addr_of_mut!(CHANNELS) };
    let reaped = channels
        .reap_exited_task(owner, scheduler)
        .unwrap_or_else(|_| fail("channel exit reap"));
    let mut revoked = 0;
    for id in reaped.ids[..reaped.count].iter().copied() {
        revoked += table.revoke_channel(id);
        if channels.queued(id) != Err(ipc::Error::InvalidChannel) {
            fail("dead channel remained queryable");
        }
    }
    let closed = table.close_all(owner);
    let expected = if count == 6 {
        (1, 1, 0, 1, 0, 1)
    } else {
        (0, 0, 0, 0, 2, 0)
    };
    if (
        reaped.count,
        reaped.discarded_messages,
        reaped.woken_peers,
        revoked,
        closed,
        table.len(),
    ) != expected
        || channels.live_channels() != 0
    {
        fail("exit resource accounting");
    }
    write_str("TURKOS/USER-HANDLE-DEATH-REAP: INDEX=0x");
    write_hex(count);
    write_str(" TASK=");
    write_str(if count == 6 { "A" } else { "B" });
    write_str(" CHANNELS=0x");
    write_hex(reaped.count);
    write_str(" DROPPED=0x");
    write_hex(reaped.discarded_messages);
    write_str(" WOKEN=0x");
    write_hex(reaped.woken_peers);
    write_str(" REVOKED=0x");
    write_hex(revoked);
    write_str(" CLOSED=0x");
    write_hex(closed);
    write_str(" LIVE=0x");
    write_hex(channels.live_channels());
    write_str(" OPEN=0x");
    write_hex(table.len());
    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, expected_return, expected_len) = 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"));
    let scheduler = unsafe { &mut *core::ptr::addr_of_mut!(user_task::SCHEDULER) };
    if 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
        || (current == 0 && count >= 1 && frame.x[22] != 1)
        || (current == 1 && count >= 5 && frame.x[22] != 2)
        || (current == 1 && count >= 10 && frame.x[23] != 3)
        || 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
        || (count == 1 && (call.args[1] != mmu::USER_DATA_VA + 16 || call.args[2] != 4))
    {
        fail("call order/arguments");
    }
    let table = unsafe { &mut *core::ptr::addr_of_mut!(HANDLES) };
    let mut exited_next = None;
    let result = if count == 0 {
        let id = unsafe { read_volatile(core::ptr::addr_of!(CHANNEL_ID)) };
        if id != 1 {
            fail("channel ID");
        }
        Some(
            table
                .open(owner, Resource::Channel(id))
                .unwrap_or_else(|_| fail("open channel")) as i64,
        )
    } else if count == 1 {
        let Ok(Resource::Channel(id)) = table.query(owner, call.args[0]) else {
            fail("send without channel authority");
        };
        let mut bytes = [0u8; 4];
        if unsafe { page_for(current).copy_from(call.args[1], &mut bytes) } != Ok(4)
            || &bytes != b"ABCD"
        {
            fail("user source copy");
        }
        let channels = unsafe { &mut *core::ptr::addr_of_mut!(CHANNELS) };
        channels
            .send(id, owner, &bytes)
            .unwrap_or_else(|_| fail("send channel"));
        if channels.queued(id) != Ok(1) {
            fail("message not queued");
        }
        write_line("TURKOS/USER-HANDLE-DEATH-IPC: A/1 SEND=ABCD QUEUED=0x1");
        Some(4)
    } else {
        match syscall::dispatch(call) {
            Outcome::HandleOpen(kind) => {
                let resource =
                    Resource::from_number(kind).unwrap_or_else(|| fail("invalid open kind"));
                Some(
                    table
                        .open(owner, resource)
                        .unwrap_or_else(|_| fail("open clock/serial")) as i64,
                )
            }
            Outcome::HandleQuery(id) => Some(
                table
                    .query(owner, id)
                    .map(|resource| resource.number() as i64)
                    .unwrap_or(syscall::EBADF),
            ),
            Outcome::Yield => Some(0),
            Outcome::Exit(code) if code == token(current) => {
                exited_next = Some(
                    scheduler
                        .exit_current(code as i32)
                        .unwrap_or_else(|_| fail("exit transition")),
                );
                exit_cleanup(owner, count, table, scheduler);
                None
            }
            _ => fail("dispatch or exit code"),
        }
    };
    if result != expected_return || table.len() != expected_len {
        fail("result/table accounting");
    }
    write_str("TURKOS/USER-HANDLE-DEATH-CALL: INDEX=0x");
    write_hex(count);
    write_str(" ID=");
    write_str(if current == 0 { "A" } else { "B" });
    write_str(" NUMBER=0x");
    write_hex(number as usize);
    write_str(" ARG=0x");
    write_hex(arg as usize);
    write_str(" RESULT=0x");
    write_hex(result.unwrap_or(token(current) as i64) as usize);
    write_str(" TABLE=0x");
    write_hex(table.len());
    write_str(" ROOT=0x");
    write_hex(root as usize);
    write_str(" ELR=0x");
    write_hex(frame.elr as usize);
    write_str("\r\n");
    unsafe { write_volatile(core::ptr::addr_of_mut!(CALLS), count + 1) };
    if number != syscall::YIELD && number != syscall::EXIT {
        frame.x[0] = result.unwrap() as u64;
        write_str("TURKOS/USER-HANDLE-DEATH-RET: INDEX=0x");
        write_hex(count);
        write_str(" X0=0x");
        write_hex(frame.x[0] as usize);
        write_str(" NEXT=0x");
        write_hex(frame.elr as usize);
        write_str("\r\n");
        return 1;
    }
    let next = if number == syscall::YIELD {
        frame.x[0] = 0;
        unsafe {
            core::ptr::write(
                core::ptr::addr_of_mut!(user_task::SAVED)
                    .cast::<irq::IrqContext>()
                    .add(current),
                *frame,
            )
        };
        scheduler
            .yield_current(false)
            .unwrap_or_else(|_| fail("yield transition"))
    } else {
        exited_next.unwrap_or_else(|| fail("exit selection absent"))
    };
    let next_index = match count {
        2 | 6 => Some(1),
        5 => Some(0),
        11 => None,
        _ => fail("scheduler call index"),
    };
    if next != next_index.map(|index| ids[index].unwrap())
        || scheduler.ready_len() != if count < 6 { 1 } else { 0 }
    {
        fail("scheduler selection/ready queue");
    }
    write_str("TURKOS/USER-HANDLE-DEATH-SWITCH: INDEX=0x");
    write_hex(count);
    write_str(" FROM=");
    write_str(if current == 0 { "A" } else { "B" });
    write_str(" TO=");
    write_str(match next_index {
        Some(0) => "A",
        Some(1) => "B",
        _ => "NONE",
    });
    write_str(" READY=0x");
    write_hex(scheduler.ready_len());
    write_str("\r\n");
    if let Some(next_index) = next_index {
        if count == 2 {
            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;
        } else {
            *frame = unsafe {
                core::ptr::read(
                    core::ptr::addr_of!(user_task::SAVED)
                        .cast::<irq::IrqContext>()
                        .add(next_index),
                )
            };
        }
        user_task::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: user_task::TOKEN_A as i32,
            })
        || scheduler.state(ids[1].unwrap())
            != Ok(State::Exited {
                code: user_task::TOKEN_B as i32,
            })
        || table.len() != 0
        || unsafe { &*core::ptr::addr_of!(CHANNELS) }.live_channels() != 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 owners/resources");
    }
    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-HANDLE-DEATH-RECLAIM: KERNEL=0x");
    write_hex(roots[0] as usize);
    write_str(" RELEASED=0x");
    write_hex((after - before) as usize);
    write_str(" OPEN=0x");
    write_hex(table.len());
    write_str(" LIVE=0x");
    write_hex(unsafe { &*core::ptr::addr_of!(CHANNELS) }.live_channels());
    write_str("\r\n");
    write_line("TURKOS/USER-HANDLE-DEATH: PASS");
    0
}