#![no_std]
#![no_main]

use core::sync::atomic::{AtomicU64, Ordering};
use lancer_user::net;
use lancer_user::show;
use lancer_user::syscall;

const TAG_NOTIFICATION: u64 = 2;
const TAG_SCHED_CONTEXT: u64 = 5;
const TAG_FRAME: u64 = 11;
const NOTIFY_BIT_DEATH: u64 = 1;
const RIGHTS_ALL: u64 = 0xFFFF;

const SLOT_UNTYPED: u64 = 30;
const SLOT_DEATH_NOTIF: u64 = 31;
const SLOTS_PER_THREAD: u64 = 3;

const DEFAULT_N: usize = 4;
const STACK_PAGE_SIZE: u64 = 4096;
const STACK_STRIDE: u64 = 8192;
const STACK_BASE: u64 = 0x0000_7000_0000_0000;

const CHURN_ITERATIONS: u64 = 10000;

const ROOT_SIZE_BITS: u64 = 10;
const SLOTS_PER_CNODE: u64 = 1 << ROOT_SIZE_BITS;
const FLAT_FIRST_SLOT: u64 = 32;
const EXPANDED_FIRST_SLOT: u64 = SLOTS_PER_CNODE;
const FLAT_MAX_THREADS: usize = ((SLOTS_PER_CNODE - FLAT_FIRST_SLOT) / SLOTS_PER_THREAD) as usize;

const BOOT_CAP_RANGE: u64 = 31;

static SHARED_COUNTER: AtomicU64 = AtomicU64::new(0);
static FIRST_THREAD_SLOT: AtomicU64 = AtomicU64::new(FLAT_FIRST_SLOT);
static THREAD_CAPACITY: AtomicU64 = AtomicU64::new(FLAT_MAX_THREADS as u64);

#[unsafe(no_mangle)]
pub extern "C" fn thread_entry(_thread_id: u64) -> ! {
    SHARED_COUNTER.fetch_add(1, Ordering::SeqCst);
    syscall::exit()
}

fn max_threads() -> usize {
    THREAD_CAPACITY.load(Ordering::Relaxed) as usize
}

fn thread_slots(i: usize) -> (u64, u64, u64) {
    let base = FIRST_THREAD_SLOT.load(Ordering::Relaxed) + (i as u64) * SLOTS_PER_THREAD;
    (base, base + 1, base + 2)
}

fn expand_cspace(child_cnode_count: u64) {
    const STAGING: u64 = 900;

    let r = syscall::cnode_create(ROOT_SIZE_BITS, STAGING);
    assert!(r >= 0, "cnode_create staging failed: {r}");

    (0..BOOT_CAP_RANGE).for_each(|slot| {
        let _ = syscall::cnode_copy_to(slot, STAGING, slot, RIGHTS_ALL);
    });

    (0..BOOT_CAP_RANGE).for_each(|slot| {
        let _ = syscall::cnode_delete(slot);
    });

    let r = syscall::cnode_copy(STAGING, 0, RIGHTS_ALL);
    assert!(r >= 0, "cnode_copy staging->0 failed: {r}");
    syscall::cnode_delete(STAGING);

    (1..=child_cnode_count).for_each(|i| {
        let r = syscall::cnode_create(ROOT_SIZE_BITS, i);
        assert!(r >= 0, "cnode_create child {i} failed: {r}");
    });

    let new_guard_bits: u64 = 64 - ROOT_SIZE_BITS * 2;
    let r = syscall::cspace_set_guard(new_guard_bits);
    assert!(r >= 0, "cspace_set_guard failed: {r}");

    let total_thread_slots = child_cnode_count * SLOTS_PER_CNODE;
    FIRST_THREAD_SLOT.store(EXPANDED_FIRST_SLOT, Ordering::Relaxed);
    THREAD_CAPACITY.store(total_thread_slots / SLOTS_PER_THREAD, Ordering::Relaxed);
}

fn ensure_capacity(needed: usize) {
    match needed <= FLAT_MAX_THREADS {
        true => {}
        false => {
            let thread_slots_needed = needed as u64 * SLOTS_PER_THREAD;
            let child_cnodes = thread_slots_needed.div_ceil(SLOTS_PER_CNODE);
            show!(thread_test, "expanding cspace: {} child cnodes for {} threads", child_cnodes, needed);
            expand_cspace(child_cnodes);
            show!(thread_test, "cspace expanded, capacity: {} threads", max_threads());
        }
    }
}

#[derive(Clone, Copy, PartialEq)]
enum Mode {
    Normal(usize),
    Stress,
}

fn parse_args(args: &[u8]) -> Mode {
    let trimmed = &args[..args.iter().position(|&b| b == 0).unwrap_or(args.len())];
    match trimmed.windows(6).any(|w| w == b"stress") {
        true => Mode::Stress,
        false => Mode::Normal(parse_n(trimmed)),
    }
}

fn parse_n(args: &[u8]) -> usize {
    let digits_start = match (0..args.len())
        .rev()
        .find(|&i| args[i] == b' ' || args[i] == b'/')
    {
        Some(pos) => pos + 1,
        None => return DEFAULT_N,
    };

    let digits = &args[digits_start..];
    match digits.is_empty() || !digits[0].is_ascii_digit() {
        true => DEFAULT_N,
        false => {
            let val = digits
                .iter()
                .take_while(|b| b.is_ascii_digit())
                .fold(0usize, |acc, &b| {
                    acc.saturating_mul(10).saturating_add((b - b'0') as usize)
                });
            match val {
                0 => DEFAULT_N,
                _ => val,
            }
        }
    }
}

fn ok(r: i64) -> Result<u64, ()> {
    match r < 0 {
        true => Err(()),
        false => Ok(r as u64),
    }
}

fn spawn_thread(i: usize) -> Result<u64, ()> {
    let (mem_slot, proc_slot, sched_slot) = thread_slots(i);
    let stack_vaddr = STACK_BASE + (i as u64) * STACK_STRIDE;

    ok(syscall::untyped_retype(SLOT_UNTYPED, TAG_FRAME, 0, mem_slot, 1))?;

    let result = spawn_thread_inner(i, mem_slot, proc_slot, sched_slot, stack_vaddr);
    if result.is_err() {
        destroy_thread(i);
    }
    result
}

fn spawn_thread_inner(
    _i: usize,
    mem_slot: u64,
    proc_slot: u64,
    sched_slot: u64,
    stack_vaddr: u64,
) -> Result<u64, ()> {
    ok(syscall::frame_map(mem_slot, stack_vaddr, 1))?;

    let stack_top = stack_vaddr + STACK_PAGE_SIZE;
    let child_pid = ok(syscall::thread_create(
        SLOT_UNTYPED, proc_slot, thread_entry as *const () as u64, stack_top, _i as u64,
    ))?;

    syscall::proc_bind_death_notif(proc_slot, SLOT_DEATH_NOTIF, NOTIFY_BIT_DEATH);

    ok(syscall::untyped_retype(SLOT_UNTYPED, TAG_SCHED_CONTEXT, 0, sched_slot, 1))?;
    ok(syscall::sched_configure(sched_slot, 5_000_000, 10_000_000, 100))?;
    ok(syscall::sched_attach(sched_slot, child_pid))?;
    ok(syscall::proc_start(proc_slot, _i as u64))?;

    Ok(child_pid)
}

fn wait_all(target: u64) {
    core::iter::repeat(())
        .take_while(|()| SHARED_COUNTER.load(Ordering::SeqCst) < target)
        .for_each(|()| {
            syscall::notify_wait(SLOT_DEATH_NOTIF);
        });
}

fn destroy_thread(i: usize) {
    let (mem_slot, proc_slot, sched_slot) = thread_slots(i);
    syscall::proc_destroy(proc_slot);
    syscall::cap_revoke(sched_slot);
    syscall::cap_revoke(mem_slot);
}

fn destroy_threads(count: u64) {
    (0..count).for_each(|i| destroy_thread(i as usize));
}

fn create_death_notif() {
    let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_NOTIFICATION, 0, SLOT_DEATH_NOTIF, 1);
    assert!(r >= 0, "untyped_retype death notif failed");
}

fn query_debug_info() -> Option<(u64, u64, u64, u64)> {
    let (rax, rsi, rdx, r8) = syscall::debug_info();
    match (rax as i64) < 0 {
        true => None,
        false => Some((rax, rsi, rdx, r8)),
    }
}

fn run_normal(n: usize) {
    ensure_capacity(n);
    let max = max_threads();
    let n = n.min(max);

    show!(thread_test, "spawning {n} threads (cap slots: {max} max)");

    SHARED_COUNTER.store(0, Ordering::SeqCst);

    create_death_notif();

    let spawned = (0..n)
        .try_fold(0u64, |count, i| match spawn_thread(i) {
            Ok(_) => Ok(count + 1),
            Err(()) => Err(count),
        })
        .unwrap_or_else(|count| count);

    show!(thread_test, "{spawned} threads started, waiting");

    wait_all(spawned);

    let val = SHARED_COUNTER.load(Ordering::SeqCst);

    destroy_threads(spawned);
    syscall::cap_revoke(SLOT_DEATH_NOTIF);

    match val == spawned {
        true => show!(thread_test, "results: {val}/{spawned} ok"),
        false => show!(thread_test, error, "results: {val}/{spawned} FAILED"),
    };
}

fn run_stress() {
    show!(thread_test, "stress mode starting");

    ensure_capacity(10000);

    let before = query_debug_info();

    let start_ms = syscall::clock_monotonic_ms();

    let max_concurrent = stress_phase1_concurrent();

    let churn_count = stress_phase2_churn();

    let elapsed_ms = syscall::clock_monotonic_ms() - start_ms;

    let after = query_debug_info();

    show!(thread_test, "--- stress results ---");
    show!(thread_test, "max concurrent:  {}", max_concurrent);
    show!(thread_test, "churn completed: {}", churn_count);
    show!(thread_test, "elapsed:         {} ms", elapsed_ms);
    match (before, after) {
        (Some((pb, ob, fb, _)), Some((pa, oa, fa, tc))) => {
            show!(thread_test, "before: pids={} objs={} frames={}", pb, ob, fb);
            show!(thread_test, "after:  pids={} objs={} frames={}", pa, oa, fa);
            show!(thread_test, "last tick cycles: {}", tc);
            let frame_delta = (fb as i64) - (fa as i64);
            match frame_delta.abs() <= 2 {
                true => show!(thread_test, "frame leak check: OK (delta={})", frame_delta),
                false => show!(
                    thread_test,
                    error,
                    "frame leak check: LEAK delta={}",
                    frame_delta
                ),
            };
        }
        _ => show!(
            thread_test,
            "(debug_info unavailable -- test kernel required)"
        ),
    }
}

fn stress_phase1_concurrent() -> u64 {
    let max = max_threads();
    SHARED_COUNTER.store(0, Ordering::SeqCst);

    create_death_notif();

    let spawned = (0..max).take_while(|&i| spawn_thread(i).is_ok()).count() as u64;

    show!(thread_test, "phase1: spawned {} / {} max", spawned, max);

    wait_all(spawned);
    destroy_threads(spawned);
    syscall::cap_revoke(SLOT_DEATH_NOTIF);

    spawned
}

fn stress_phase2_churn() -> u64 {
    SHARED_COUNTER.store(0, Ordering::SeqCst);
    let mut completed = 0u64;

    (0..CHURN_ITERATIONS).for_each(|iter| {
        create_death_notif();
        match spawn_thread(0) {
            Ok(_) => {
                core::iter::repeat(())
                    .take_while(|()| SHARED_COUNTER.load(Ordering::SeqCst) <= completed)
                    .for_each(|()| {
                        syscall::notify_wait(SLOT_DEATH_NOTIF);
                    });
                completed += 1;

                destroy_thread(0);
            }
            Err(()) => {
                show!(thread_test, "phase2: churn iter {} spawn failed", iter);
            }
        }
        syscall::cap_revoke(SLOT_DEATH_NOTIF);
    });

    show!(
        thread_test,
        "phase2: {} sequential churn iterations",
        completed
    );
    completed
}

#[unsafe(no_mangle)]
pub extern "C" fn lancer_main() -> ! {
    let mode = match net::init() {
        Some((rx, _tx)) => {
            let mut buf = [0u8; 64];
            let len = net::recv_args(&rx, &mut buf);
            parse_args(&buf[..len])
        }
        None => Mode::Normal(DEFAULT_N),
    };

    match mode {
        Mode::Normal(n) => run_normal(n),
        Mode::Stress => run_stress(),
    };

    syscall::send(1, [1, 0, 0, 0, 0, 0]);
    syscall::exit()
}