//! thread-test — danos's multi-threaded exerciser (docs/threading-plan.md M2, M3).
//!
//! Two modes, chosen by argv[1] (default "spawn"):
//!   spawn — M2: one worker writes a shared global; the main thread observes it, proving
//!           `Thread.spawn` started a task in the **same** address space.
//!   join  — M3: N workers each do K atomic increments on a shared counter and stamp the
//!           core they ran on; the main thread `join`s all N and checks the total is
//!           exactly N*K (every worker ran, join waited) and that >1 core was used
//!           (genuine parallelism). Then a detached worker proves `detach` runs and
//!           needs no join.
//!
//! Built multi-threaded (`addThreadedUserBinary`) so atomics/shared reads are real.

const std = @import("std");
const process = @import("process");
const time = @import("time");
const memory = @import("memory");
const logging = @import("logging");
const Thread = @import("thread").Thread;

fn write(comptime s: []const u8) void {
    _ = logging.write(s);
}

// --- M2: spawn mode ---------------------------------------------------------

var shared_value: u32 = 0;
var spawn_done = std.atomic.Value(u32).init(0);
const sentinel: u32 = 0xA5A5;

fn spawnWorker() void {
    shared_value = sentinel;
    spawn_done.store(1, .release);
}

fn runSpawnMode() void {
    write("thread-test: starting\n");
    _ = Thread.spawn(.{}, spawnWorker, .{}) catch {
        write("thread-test: FAIL spawn refused\n");
        return;
    };
    var spins: usize = 0;
    while (spawn_done.load(.acquire) == 0 and spins < 50_000_000) : (spins += 1) {
        process.yield();
    }
    if (spawn_done.load(.acquire) == 1 and shared_value == sentinel) {
        write("thread-test: child ran in shared address space ok\n");
    } else {
        write("thread-test: FAIL worker did not update shared memory\n");
    }
}

// --- M3: join mode ----------------------------------------------------------

const worker_count: u32 = 4;
const iterations: u64 = 100_000;

var counter = std.atomic.Value(u64).init(0);
var cores_seen = std.atomic.Value(u32).init(0);

fn joinWorker() void {
    var i: u64 = 0;
    while (i < iterations) : (i += 1) {
        _ = counter.fetchAdd(1, .monotonic);
        if (i % 1000 == 0) stampCore(); // periodic: catches cross-core migration too
    }
    stampCore();
}

fn stampCore() void {
    const core = Thread.currentCore();
    if (core < 32) _ = cores_seen.fetchOr(@as(u32, 1) << @intCast(core), .monotonic);
}

var detach_done = std.atomic.Value(u32).init(0);

fn detachWorker() void {
    detach_done.store(1, .release);
}

fn noopWorker() void {}

fn runJoinMode() void {
    write("thread-test: join mode starting\n");

    var threads: [worker_count]Thread = undefined;
    var spawned: u32 = 0;
    while (spawned < worker_count) : (spawned += 1) {
        threads[spawned] = Thread.spawn(.{}, joinWorker, .{}) catch break;
    }
    if (spawned != worker_count) {
        write("thread-test: FAIL could not spawn all workers\n");
        return;
    }
    for (threads[0..spawned]) |t| t.join();

    const total = counter.load(.acquire);
    const cores = @popCount(cores_seen.load(.acquire));
    if (total != worker_count * iterations) {
        write("thread-test: FAIL counter mismatch (a worker was lost or join did not wait)\n");
        return;
    }
    if (cores <= 1) {
        write("thread-test: FAIL workers never ran on more than one core\n");
        return;
    }

    // detach: the worker runs and we never join it.
    const dt = Thread.spawn(.{}, detachWorker, .{}) catch {
        write("thread-test: FAIL detach spawn refused\n");
        return;
    };
    dt.detach();
    var spins: usize = 0;
    while (detach_done.load(.acquire) == 0 and spins < 50_000_000) : (spins += 1) {
        process.yield();
    }
    if (detach_done.load(.acquire) != 1) {
        write("thread-test: FAIL detached worker did not run\n");
        return;
    }

    // Prove join needs no per-thread kernel endpoint (M9): many spawn+join cycles. Under
    // the old per-thread-endpoint scheme these leaked handles and would exhaust the
    // 16-slot handle table well before 40; here they all succeed.
    var cycle: u32 = 0;
    while (cycle < 40) : (cycle += 1) {
        const th = Thread.spawn(.{}, noopWorker, .{}) catch {
            write("thread-test: FAIL spawn exhausted across join cycles (endpoint leak?)\n");
            return;
        };
        th.join();
    }

    write("thread-test: join ok\n"); // the M3/M9 verdict marker
}

// --- M4: futex mode ---------------------------------------------------------

const Futex = Thread.Futex;

var futex_word = std.atomic.Value(u32).init(0);
var waiter_parked = std.atomic.Value(u32).init(0);

fn futexWaiter() void {
    write("thread-futex: waiting\n");
    waiter_parked.store(1, .release);
    // Block while the word is still 0; the waker sets it to 1 and wakes us.
    while (futex_word.load(.acquire) == 0) {
        Futex.wait(&futex_word, 0);
    }
    write("thread-futex: woke\n");
}

fn runFutexMode() void {
    write("thread-futex: starting\n");

    const waiter = Thread.spawn(.{}, futexWaiter, .{}) catch {
        write("thread-futex: FAIL spawn refused\n");
        return;
    };
    // Let the waiter reach its wait, then give it a beat to actually park in-kernel.
    var spins: usize = 0;
    while (waiter_parked.load(.acquire) == 0 and spins < 50_000_000) : (spins += 1) {
        process.yield();
    }
    time.sleepMillis(50);

    // The handshake: publish the value, then wake the parked waiter.
    futex_word.store(1, .release);
    write("thread-futex: waking\n");
    Futex.wake(&futex_word, 1);

    waiter.join(); // returns once the waiter woke and printed "woke"

    // Timeout: nobody ever wakes this word, so timedWait must report a timeout.
    var lonely = std.atomic.Value(u32).init(0);
    if (Futex.timedWait(&lonely, 0, 100_000_000)) |_| {
        write("thread-futex: FAIL timedWait did not time out\n");
        return;
    } else |_| {}
    write("thread-futex: timeout ok\n");

    write("thread-futex: ok\n"); // the M4 verdict marker
}

// --- M5: mutex mode (bounded producer/consumer over Mutex + Condition) ------

const Mutex = Thread.Mutex;
const Condition = Thread.Condition;

const producers: u32 = 2;
const consumers: u32 = 2;
const per_producer: u32 = 1000;
const per_consumer: u32 = 1000; // producers*per_producer == consumers*per_consumer (balanced)
const total_items: u32 = producers * per_producer;
const ring_cap: usize = 8; // small, so producers block on full and consumers on empty

var ring: [ring_cap]u32 = undefined;
var ring_count: usize = 0;
var ring_head: usize = 0;
var ring_tail: usize = 0;

var pc_mutex = Mutex{};
var not_full = Condition{};
var not_empty = Condition{};

// Verified outside the lock: the checksum and tally of everything consumed.
var consumed_sum = std.atomic.Value(u64).init(0);
var consumed_count = std.atomic.Value(u32).init(0);

fn producer(base: u32) void {
    var i: u32 = 0;
    while (i < per_producer) : (i += 1) {
        const item = base + i;
        pc_mutex.lock();
        while (ring_count == ring_cap) not_full.wait(&pc_mutex);
        ring[ring_tail] = item;
        ring_tail = (ring_tail + 1) % ring_cap;
        ring_count += 1;
        pc_mutex.unlock();
        not_empty.signal();
    }
}

fn consumer() void {
    var i: u32 = 0;
    while (i < per_consumer) : (i += 1) {
        pc_mutex.lock();
        while (ring_count == 0) not_empty.wait(&pc_mutex);
        const item = ring[ring_head];
        ring_head = (ring_head + 1) % ring_cap;
        ring_count -= 1;
        pc_mutex.unlock();
        not_full.signal();
        _ = consumed_sum.fetchAdd(item, .monotonic);
        _ = consumed_count.fetchAdd(1, .monotonic);
    }
}

fn runMutexMode() void {
    write("thread-mutex: starting\n");

    var threads: [producers + consumers]Thread = undefined;
    var n: usize = 0;
    var p: u32 = 0;
    while (p < producers) : (p += 1) {
        threads[n] = Thread.spawn(.{}, producer, .{p * per_producer}) catch {
            write("thread-mutex: FAIL producer spawn\n");
            return;
        };
        n += 1;
    }
    var c: u32 = 0;
    while (c < consumers) : (c += 1) {
        threads[n] = Thread.spawn(.{}, consumer, .{}) catch {
            write("thread-mutex: FAIL consumer spawn\n");
            return;
        };
        n += 1;
    }
    for (threads[0..n]) |t| t.join();

    // Every item 0..total_items-1 was produced exactly once; if the mutex/condition are
    // correct, each was consumed exactly once, so the checksum matches.
    const expected_sum: u64 = @as(u64, total_items) * (total_items - 1) / 2;
    if (consumed_count.load(.acquire) != total_items) {
        write("thread-mutex: FAIL wrong number of items consumed\n");
        return;
    }
    if (consumed_sum.load(.acquire) != expected_sum) {
        write("thread-mutex: FAIL checksum mismatch (item lost or duplicated)\n");
        return;
    }
    write("thread-mutex: ok\n"); // the M5 verdict marker
}

// --- M6: id mode (getCurrentId identity) ------------------------------------

var worker_ids: [2]std.atomic.Value(u32) = .{ std.atomic.Value(u32).init(0), std.atomic.Value(u32).init(0) };

fn idWorker(slot: usize) void {
    worker_ids[slot].store(Thread.getCurrentId(), .release);
}

fn runIdMode() void {
    write("thread-id: starting\n");
    const main_id = Thread.getCurrentId();

    const t0 = Thread.spawn(.{}, idWorker, .{@as(usize, 0)}) catch {
        write("thread-id: FAIL spawn\n");
        return;
    };
    const t1 = Thread.spawn(.{}, idWorker, .{@as(usize, 1)}) catch {
        write("thread-id: FAIL spawn\n");
        return;
    };
    t0.join();
    t1.join();

    const id0 = worker_ids[0].load(.acquire);
    const id1 = worker_ids[1].load(.acquire);
    // Each thread has its own kernel task id: all three distinct and non-zero.
    if (main_id == 0 or id0 == 0 or id1 == 0) {
        write("thread-id: FAIL a thread reported id 0\n");
        return;
    }
    if (id0 == id1 or id0 == main_id or id1 == main_id) {
        write("thread-id: FAIL thread ids collided\n");
        return;
    }
    write("thread-id: ok\n"); // the M6 verdict marker
}

// --- M7: alloc mode (concurrent heap allocation) ----------------------------

const alloc_threads: u32 = 4;
const allocs_per_thread: u32 = 500;
var allocs_clean = std.atomic.Value(u32).init(0);

fn allocWorker(seed: u32) void {
    const gpa = memory.allocator();
    var rng: u32 = seed | 1;
    var round: u32 = 0;
    while (round < allocs_per_thread) : (round += 1) {
        rng = rng *% 1664525 +% 1013904223; // cheap LCG for varied sizes
        const size: usize = 16 + (rng % 4080); // 16..4095 bytes
        const buf = gpa.alloc(u8, size) catch return; // OOM: don't count this thread clean
        const pattern: u8 = @truncate(seed +% round);
        @memset(buf, pattern);
        // Nothing else should touch our block; if a concurrent allocation overlapped it,
        // one of us would read the other's pattern here.
        var ok = true;
        for (buf) |b| {
            if (b != pattern) ok = false;
        }
        gpa.free(buf);
        if (!ok) return; // corruption — leave without counting clean
    }
    _ = allocs_clean.fetchAdd(1, .monotonic);
}

fn runAllocMode() void {
    write("thread-alloc: starting\n");
    var threads: [alloc_threads]Thread = undefined;
    var n: u32 = 0;
    while (n < alloc_threads) : (n += 1) {
        threads[n] = Thread.spawn(.{}, allocWorker, .{n +% 1}) catch {
            write("thread-alloc: FAIL spawn\n");
            return;
        };
    }
    for (threads[0..alloc_threads]) |t| t.join();

    // Every thread must have completed all rounds with each block intact — proof the
    // shared heap and the per-address_space mmap arena are safe under concurrent allocation.
    if (allocs_clean.load(.acquire) != alloc_threads) {
        write("thread-alloc: FAIL corruption or OOM under concurrent allocation\n");
        return;
    }
    write("thread-alloc: ok\n"); // the M7 verdict marker
}

// --- M10: tls mode (per-thread FS base storage) -----------------------------

fn writeTlsSlot(value: u64) void {
    asm volatile ("movq %[v], %%fs:8"
        :
        : [v] "r" (value),
        : .{ .memory = true });
}

fn readTlsSlot() u64 {
    return asm volatile ("movq %%fs:8, %[out]"
        : [out] "=r" (-> u64),
        :
        : .{ .memory = true });
}

var tls_written = std.atomic.Value(u32).init(0);
var tls_ok = std.atomic.Value(u32).init(0);

fn tlsWorker(marker: u64) void {
    writeTlsSlot(marker);
    _ = tls_written.fetchAdd(1, .release);
    // Wait until both threads have written their own slot. If the FS base were shared, the
    // second write would clobber the first, and the read below would return the wrong
    // marker — cross-talk. A per-thread FS base keeps each thread's slot private.
    var spins: usize = 0;
    while (tls_written.load(.acquire) < 2 and spins < 50_000_000) : (spins += 1) {
        process.yield();
    }
    if (readTlsSlot() == marker and Thread.getCurrentId() != 0) {
        _ = tls_ok.fetchAdd(1, .monotonic);
    }
}

fn runTlsMode() void {
    write("thread-tls: starting\n");
    const t0 = Thread.spawn(.{}, tlsWorker, .{@as(u64, 0xAAAA_0000)}) catch {
        write("thread-tls: FAIL spawn\n");
        return;
    };
    const t1 = Thread.spawn(.{}, tlsWorker, .{@as(u64, 0xBBBB_0000)}) catch {
        write("thread-tls: FAIL spawn\n");
        return;
    };
    t0.join();
    t1.join();
    if (tls_ok.load(.acquire) == 2) {
        write("thread-tls: ok\n"); // the M10 verdict marker
    } else {
        write("thread-tls: FAIL cross-talk (FS base not per-thread)\n");
    }
}

// --- M11: rwlock mode (readers/writers over an RwLock) ----------------------

const RwLock = Thread.RwLock;

var rwlock = RwLock{};
var rw_a: u64 = 0;
var rw_b: u64 = 0; // invariant while any lock is held: rw_a == rw_b
var rw_stop = std.atomic.Value(u32).init(0);
var rw_violations = std.atomic.Value(u32).init(0);
var rw_reads = std.atomic.Value(u64).init(0);

fn rwWriter() void {
    var v: u64 = 1;
    while (rw_stop.load(.acquire) == 0) : (v +%= 1) {
        rwlock.lock(); // exclusive: no reader may observe the gap between the two writes
        rw_a = v;
        rw_b = v;
        rwlock.unlock();
    }
}

fn rwReader() void {
    const reads: u64 = 50_000;
    var i: u64 = 0;
    while (i < reads) : (i += 1) {
        rwlock.lockShared();
        if (rw_a != rw_b) _ = rw_violations.fetchAdd(1, .monotonic); // saw a half-write!
        rwlock.unlockShared();
    }
    _ = rw_reads.fetchAdd(reads, .monotonic);
}

fn runRwlockMode() void {
    write("thread-rwlock: starting\n");
    var writers: [2]Thread = undefined;
    var readers: [3]Thread = undefined;
    for (&writers) |*w| {
        w.* = Thread.spawn(.{}, rwWriter, .{}) catch {
            write("thread-rwlock: FAIL spawn\n");
            return;
        };
    }
    for (&readers) |*r| {
        r.* = Thread.spawn(.{}, rwReader, .{}) catch {
            write("thread-rwlock: FAIL spawn\n");
            return;
        };
    }
    for (readers) |r| r.join();
    rw_stop.store(1, .release); // readers done → stop the writers
    for (writers) |w| w.join();

    if (rw_violations.load(.acquire) == 0 and rw_reads.load(.acquire) > 0) {
        write("thread-rwlock: ok\n"); // the M11 verdict marker
    } else {
        write("thread-rwlock: FAIL reader observed a half-written value\n");
    }
}

// --- shared-fate modes (docs/shared-fate-plan.md M4) -------------------------

fn faultNullPage() void {
    @as(*volatile u32, @ptrFromInt(0x10)).* = 1; // unmapped null page: #PF, group death
}

fn faultingWorker() void {
    write("thread-test: worker faulting\n");
    faultNullPage();
}

/// fault-worker: a worker faults; shared fate must end the whole group, so the
/// main thread parks forever and never prints anything more.
fn runFaultWorkerMode() void {
    write("thread-test: spawning faulting worker\n");
    _ = Thread.spawn(.{}, faultingWorker, .{}) catch {
        write("thread-test: FAIL spawn refused\n");
        return;
    };
    while (true) process.yield();
}

fn spinningWorker() void {
    while (true) {} // no system calls: only a tick can deliver a deferred kill
}

/// spin-forever: a kill target. The worker spins without syscalls (the condemned
/// path); the main thread yields (the parked path).
fn runSpinForeverMode() void {
    _ = Thread.spawn(.{}, spinningWorker, .{}) catch {
        write("thread-test: FAIL spawn refused\n");
        return;
    };
    write("thread-test: spinning\n");
    while (true) process.yield();
}

fn exitingWorker() void {
    write("thread-test: worker exiting the process\n");
    process.exit(3); // exit from ANY thread is group death (.aborted)
}

/// exit-worker: a WORKER calls exit(3); the group must die with the leader's
/// reason reading .aborted.
fn runExitWorkerMode() void {
    _ = Thread.spawn(.{}, exitingWorker, .{}) catch {
        write("thread-test: FAIL spawn refused\n");
        return;
    };
    while (true) process.yield();
}

var race_go = std.atomic.Value(u32).init(0);

fn racingWorker() void {
    while (race_go.load(.acquire) == 0) {}
    faultNullPage();
}

/// race: two members fault as near-simultaneously as user space can arrange —
/// the group-dying latch must make the two triggers count as one death.
fn runRaceMode() void {
    _ = Thread.spawn(.{}, racingWorker, .{}) catch {
        write("thread-test: FAIL spawn refused\n");
        return;
    };
    write("thread-test: racing\n");
    race_go.store(1, .release);
    faultNullPage();
}

var leader_exit_done = std.atomic.Value(u32).init(0);

fn patientWorker() void {
    while (leader_exit_done.load(.acquire) == 0) process.yield();
}

/// leader-exit: the MAIN thread asks for thread_exit; the kernel must refuse
/// (-EPERM) and the worker must be entirely unaffected.
fn runLeaderExitMode() void {
    const worker = Thread.spawn(.{}, patientWorker, .{}) catch {
        write("thread-test: FAIL spawn refused\n");
        return;
    };
    Thread.tryExitCurrent(); // refused: we are the leader
    write("thread-test: leader thread_exit refused ok\n");
    leader_exit_done.store(1, .release);
    worker.join();
}

fn promptWorker() void {}

/// solo: regression — a WORKER's thread_exit stays per-thread; the sibling
/// (main) survives it.
fn runSoloMode() void {
    const worker = Thread.spawn(.{}, promptWorker, .{}) catch {
        write("thread-test: FAIL spawn refused\n");
        return;
    };
    worker.join();
    write("thread-test: solo sibling survived ok\n");
}

const shm_pattern_length: usize = 4096;
var shm_region_base = std.atomic.Value(usize).init(0);

fn patternByte(i: usize) u8 {
    return @truncate(i *% 31 +% 7);
}

fn shmWorker() void {
    const region = memory.sharedCreate(shm_pattern_length) orelse {
        write("thread-shm: FAIL create refused\n");
        return;
    };
    for (0..shm_pattern_length) |i| region.ptr[i] = patternByte(i);
    shm_region_base.store(@intFromPtr(region.ptr), .release);
    // Returning thread_exits this worker: its handle table — holding the
    // region's ONLY handle — closes. The sibling's view must survive on the
    // mapping reference (docs/shared-fate-plan.md M3).
}

/// shm-worker: the M3 use-after-free regression. The worker creates and fills a
/// shared-memory region and dies; the main thread churns the frame allocator and
/// then checks the mapping is intact — freed frames would have been reused and
/// scribbled on.
fn runShmWorkerMode() void {
    const worker = Thread.spawn(.{}, shmWorker, .{}) catch {
        write("thread-test: FAIL spawn refused\n");
        return;
    };
    worker.join();
    const base = shm_region_base.load(.acquire);
    if (base == 0) return; // the worker already printed the failure
    var churn: usize = 0;
    while (churn < 8) : (churn += 1) {
        const noise = memory.sharedCreate(shm_pattern_length) orelse break;
        @memset(noise.ptr[0..shm_pattern_length], 0xFF);
    }
    const view: [*]const u8 = @ptrFromInt(base);
    var intact = true;
    for (0..shm_pattern_length) |i| {
        if (view[i] != patternByte(i)) intact = false;
    }
    if (intact) {
        write("thread-shm: mapping survives creator ok\n");
    } else {
        write("thread-shm: FAIL mapping corrupted after creator death\n");
    }
}

pub fn main(init: process.Init) void {
    const mode = init.arguments.get(1) orelse "spawn";
    if (std.mem.eql(u8, mode, "join")) {
        runJoinMode();
    } else if (std.mem.eql(u8, mode, "futex")) {
        runFutexMode();
    } else if (std.mem.eql(u8, mode, "mutex")) {
        runMutexMode();
    } else if (std.mem.eql(u8, mode, "id")) {
        runIdMode();
    } else if (std.mem.eql(u8, mode, "alloc")) {
        runAllocMode();
    } else if (std.mem.eql(u8, mode, "tls")) {
        runTlsMode();
    } else if (std.mem.eql(u8, mode, "rwlock")) {
        runRwlockMode();
    } else if (std.mem.eql(u8, mode, "fault-worker")) {
        runFaultWorkerMode();
    } else if (std.mem.eql(u8, mode, "spin-forever")) {
        runSpinForeverMode();
    } else if (std.mem.eql(u8, mode, "exit-worker")) {
        runExitWorkerMode();
    } else if (std.mem.eql(u8, mode, "race")) {
        runRaceMode();
    } else if (std.mem.eql(u8, mode, "leader-exit")) {
        runLeaderExitMode();
    } else if (std.mem.eql(u8, mode, "solo")) {
        runSoloMode();
    } else if (std.mem.eql(u8, mode, "shm-worker")) {
        runShmWorkerMode();
    } else {
        runSpawnMode();
    }
}