//! 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
//!           `runtime.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 runtime = @import("runtime");

fn write(comptime s: []const u8) void {
    _ = runtime.system.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");
    _ = runtime.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) {
        runtime.system.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 = runtime.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]runtime.Thread = undefined;
    var spawned: u32 = 0;
    while (spawned < worker_count) : (spawned += 1) {
        threads[spawned] = runtime.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 = runtime.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) {
        runtime.system.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 = runtime.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 = runtime.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 = runtime.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) {
        runtime.system.yield();
    }
    runtime.system.sleep(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 = runtime.Thread.Mutex;
const Condition = runtime.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]runtime.Thread = undefined;
    var n: usize = 0;
    var p: u32 = 0;
    while (p < producers) : (p += 1) {
        threads[n] = runtime.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] = runtime.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(runtime.Thread.getCurrentId(), .release);
}

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

    const t0 = runtime.Thread.spawn(.{}, idWorker, .{@as(usize, 0)}) catch {
        write("thread-id: FAIL spawn\n");
        return;
    };
    const t1 = runtime.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 = runtime.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]runtime.Thread = undefined;
    var n: u32 = 0;
    while (n < alloc_threads) : (n += 1) {
        threads[n] = runtime.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) {
        runtime.system.yield();
    }
    if (readTlsSlot() == marker and runtime.Thread.getCurrentId() != 0) {
        _ = tls_ok.fetchAdd(1, .monotonic);
    }
}

fn runTlsMode() void {
    write("thread-tls: starting\n");
    const t0 = runtime.Thread.spawn(.{}, tlsWorker, .{@as(u64, 0xAAAA_0000)}) catch {
        write("thread-tls: FAIL spawn\n");
        return;
    };
    const t1 = runtime.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 = runtime.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]runtime.Thread = undefined;
    var readers: [3]runtime.Thread = undefined;
    for (&writers) |*w| {
        w.* = runtime.Thread.spawn(.{}, rwWriter, .{}) catch {
            write("thread-rwlock: FAIL spawn\n");
            return;
        };
    }
    for (&readers) |*r| {
        r.* = runtime.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");
    }
}

pub fn main(init: runtime.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 {
        runSpawnMode();
    }
}
