/*
    * thread.h
    * Userspace threading primitives for MontaukOS
    * Copyright (c) 2026 Daniel Hammer
    *
    * Threads share their parent's address space, IPC handles, cwd, and
    * heap allocator. The caller owns the thread stack: thread_spawn
    * allocates one out of the user heap and the kernel uses it verbatim.
    *
    * Lifetimes:
    *   - thread_spawn returns a positive TID on success.
    *   - thread_join blocks until the target TID exits, frees its kernel
    *     stack, then returns its exit code via *out_code.
    *   - thread_exit terminates only the current thread. If the main thread
    *     calls exit() (SYS_EXIT) the whole process tears down, killing any
    *     surviving sibling threads.
    *   - thread_self returns the calling thread's TID.
*/

#pragma once
#include <cstdint>
#include <Api/Syscall.hpp>
#include <montauk/syscall.h>
#include <montauk/heap.h>

namespace montauk {

    // Default per-thread stack size (64 KiB). Enough for typical app work;
    // matches the main thread's 32 KiB lower bound with headroom for
    // TrueType rendering call chains.
    static constexpr uint64_t DEFAULT_THREAD_STACK_BYTES = 64 * 1024;

    using ThreadEntry = int (*)(void* arg);

    // Terminate the calling thread. Never returns. If called by the main
    // thread this is equivalent to exit() (the whole process exits).
    [[noreturn]] inline void thread_exit(int code = 0) {
        syscall1(montauk::abi::SYS_THREAD_EXIT, (uint64_t)code);
        __builtin_unreachable();
    }

    namespace detail {
        struct ThreadCtx {
            ThreadEntry user_entry;
            void* user_arg;
            void* stack_base;
            int tid;
            ThreadCtx* next;
        };

        inline ThreadCtx* g_thread_records = nullptr;
        inline volatile uint32_t g_thread_records_lock = 0;

        inline void records_lock() {
            while (__atomic_exchange_n(&g_thread_records_lock, 1, __ATOMIC_ACQUIRE) != 0)
                montauk::yield();
        }

        inline void records_unlock() {
            __atomic_store_n(&g_thread_records_lock, 0, __ATOMIC_RELEASE);
        }

        // Userspace trampoline: bridges from the raw entry the kernel jumps
        // to into the typed entry, then funnels into SYS_THREAD_EXIT. We
        // route the exit through libc rather than relying on a kernel-side
        // exit stub on the user stack -- the kernel never writes to user
        // memory on this path.
        //
        // The thread's stack itself is intentionally not freed here: we are
        // still running on it. It is reclaimed by a successful thread_join,
        // or as part of whole-process teardown if the thread is never joined.
        [[noreturn]] inline void thread_trampoline(detail::ThreadCtx* ctx) {
            // A sibling CPU can start the thread before thread_spawn has
            // returned its TID. Wait until the parent has published the record
            // needed by thread_join to reclaim this stack.
            while (__atomic_load_n(&ctx->tid, __ATOMIC_ACQUIRE) == 0)
                montauk::yield();
            int code = ctx->user_entry(ctx->user_arg);
            thread_exit(code);
        }
    }

    // Spawn a new thread that begins executing `entry(arg)`. Returns the
    // new TID on success, or -1 on failure. The thread's stack is
    // allocated from the user heap. The exiting thread cannot free the stack
    // it is running on, so thread_join reclaims both it and the trampoline
    // context after the kernel has reaped the sibling.
    inline int thread_spawn(ThreadEntry entry, void* arg,
                            uint64_t stack_bytes = 0) {
        if (entry == nullptr) return -1;
        if (stack_bytes == 0) stack_bytes = DEFAULT_THREAD_STACK_BYTES;
        stack_bytes = (stack_bytes + 15) & ~15ULL;

        void* stack = montauk::malloc(stack_bytes);
        if (stack == nullptr) return -1;
        auto* ctx = (detail::ThreadCtx*)montauk::malloc(sizeof(detail::ThreadCtx));
        if (ctx == nullptr) {
            montauk::mfree(stack);
            return -1;
        }
        ctx->user_entry = entry;
        ctx->user_arg = arg;
        ctx->stack_base = stack;
        ctx->tid = 0;
        ctx->next = nullptr;

        uint64_t stack_top = ((uint64_t)stack + stack_bytes) & ~0xFULL;
        int tid = (int)syscall3(montauk::abi::SYS_THREAD_SPAWN,
                                (uint64_t)&detail::thread_trampoline,
                                (uint64_t)ctx, stack_top);
        if (tid < 0) {
            montauk::mfree(ctx);
            montauk::mfree(stack);
            return -1;
        }
        detail::records_lock();
        ctx->next = detail::g_thread_records;
        detail::g_thread_records = ctx;
        __atomic_store_n(&ctx->tid, tid, __ATOMIC_RELEASE);
        detail::records_unlock();
        return tid;
    }

    // Block until the thread identified by `tid` terminates. Returns 0 on
    // success (with the thread's exit code in *out_code if non-null) or
    // -1 if `tid` is not a joinable sibling.
    inline int thread_join(int tid, int* out_code = nullptr) {
        int result = (int)syscall2(montauk::abi::SYS_THREAD_JOIN,
                                   (uint64_t)tid, (uint64_t)out_code);
        if (result == 0) {
            detail::records_lock();
            detail::ThreadCtx** link = &detail::g_thread_records;
            while (*link != nullptr && (*link)->tid != tid)
                link = &(*link)->next;
            detail::ThreadCtx* ctx = *link;
            if (ctx != nullptr) *link = ctx->next;
            detail::records_unlock();
            if (ctx != nullptr) {
                montauk::mfree(ctx->stack_base);
                montauk::mfree(ctx);
            }
        }
        return result;
    }

    // Return the calling thread's TID (== getpid() for the main thread).
    inline int thread_self() {
        return (int)syscall0(montauk::abi::SYS_THREAD_SELF);
    }

    // Lightweight mutex backed by a single atomic word + the kernel
    // yield syscall. Adequate for short critical sections; a heavier
    // primitive can wrap mailbox_recv when blocking semantics are needed.
    struct Mutex {
        volatile uint32_t locked = 0;

        void lock() {
            while (__atomic_exchange_n(&locked, 1, __ATOMIC_ACQUIRE) != 0) {
                montauk::yield();
            }
        }

        bool try_lock() {
            return __atomic_exchange_n(&locked, 1, __ATOMIC_ACQUIRE) == 0;
        }

        void unlock() {
            __atomic_store_n(&locked, 0, __ATOMIC_RELEASE);
        }
    };

}