#ifndef SCRATCH_H
#define SCRATCH_H

#include "arena.h"
#include "sys.h"
#include "thread_local.h"

#define SCRATCH_ARENAS_PER_THREAD 2
#define SCRATCH_ARENA_RESERVE (GIGABYTE(64))

/* ========================== *
 * Thread local state
 * ========================== */

struct scratch_ctx {
    struct arena *arenas[SCRATCH_ARENAS_PER_THREAD];

#if RTC
    u64 next_scratch_id;
    u64 scratch_id_stack[16384];
    u64 scratch_id_stack_count;
#endif
};

THREAD_LOCAL_VAR_DECL_EXTERN(tl_scratch_ctx, struct scratch_ctx);

/* ========================== *
 * Scratch begin
 * ========================== */

INLINE void scratch_dbg_push(struct scratch_ctx *ctx, struct arena_temp *temp)
{
#if RTC
    if (ctx->scratch_id_stack_count >= ARRAY_COUNT(ctx->scratch_id_stack)) {
        sys_panic(LIT("Max debug scratch depth reached"));
    }
    temp->scratch_id = ctx->next_scratch_id++;
    ctx->scratch_id_stack[ctx->scratch_id_stack_count++] = temp->scratch_id;
#else
    (UNUSED)ctx;
    (UNUSED)temp;
#endif
}

/* Any parameterized arenas in the caller's scope should be passed into this
 * function as a potential "conflict". This is to prevent friction in case the
 * passed arena is itself a scratch arena from another scope (since
 * parameterized arenas are often used to allocate persistent results for the
 * caller).
 *
 * Use `scratch_begin_no_conflict` instead if there is no arena in the current
 * scope that could potentially be a scratch arena from another scope. */
#define scratch_begin(potential_conflict) _scratch_begin(potential_conflict)

INLINE struct arena_temp _scratch_begin(struct arena *potential_conflict)
{
    /* This function is currently hard-coded to support 2 scratch arenas */
    CT_ASSERT(SCRATCH_ARENAS_PER_THREAD == 2);

    /* Use `scratch_begin_no_conflict` if no conflicts are present */
    ASSERT(potential_conflict != NULL);

    struct scratch_ctx *ctx = (struct scratch_ctx *)thread_local_var_eval(&tl_scratch_ctx);
    struct arena *scratch_arena = ctx->arenas[0];
    if (potential_conflict && scratch_arena == potential_conflict) {
        scratch_arena = ctx->arenas[1];
    }
    struct arena_temp temp = arena_temp_begin(scratch_arena);
    scratch_dbg_push(ctx, &temp);
    return temp;
}

/* This macro declares an unused "arena" variable that will error if an existing "arena"
 * variable is present (due to shadowing). This is for catching obvious cases of
 * `scratch_begin_no_conflict` getting called when an `arena` variable already
 * exists in the caller's scope (`scratch_begin(arena)` should be called
 * instead). */
#define scratch_begin_no_conflict() \
    _scratch_begin_no_conflict(); \
    do { \
        u8 arena = 0; \
        (UNUSED)arena; \
    } while (0)

INLINE struct arena_temp _scratch_begin_no_conflict(void)
{
    struct scratch_ctx *ctx = (struct scratch_ctx *)thread_local_var_eval(&tl_scratch_ctx);
    struct arena *scratch_arena = ctx->arenas[0];
    struct arena_temp temp = arena_temp_begin(scratch_arena);
    scratch_dbg_push(ctx, &temp);
    return temp;
}

/* ========================== *
 * Scratch end
 * ========================== */

INLINE void scratch_end(struct arena_temp scratch_temp)
{
#if RTC
    struct scratch_ctx *ctx = (struct scratch_ctx *)thread_local_var_eval(&tl_scratch_ctx);
    if (ctx->scratch_id_stack_count > 0) {
        u64 scratch_id = scratch_temp.scratch_id;
        u64 expected_id = ctx->scratch_id_stack[--ctx->scratch_id_stack_count];
        /* This assertion exists to catch cases where a scratch_end was forgotten.
         * It will fail if a scratch arena is reset out of order.
         * E.g. there is a missing scratch_end somewhere on a different scratch
         * arena (one that was created between the scratch_begin & the
         * scratch_end of the arena being reset here). */
        ASSERT(scratch_id == expected_id);
    }
#endif
    arena_temp_end(scratch_temp);
}

#endif