#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 temp_arena *temp)
{
#if RTC
    if (ctx->scratch_id_stack_count >= ARRAY_COUNT(ctx->scratch_id_stack)) {
        sys_panic(STR("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 arena parameters in the calling function's context should be passed into this
 * function as a potential `conflict`. This is to prevent conflicts when the
 * context's arena is itself a scratch arena (since parameterized arenas are
 * often used to allocate persistent results for the caller).
 *
 * Call `scratch_begin_no_conflict` instead if there is no arena in the current
 * context that could potentially be a scratch arena. */
#define scratch_begin(c) _scratch_begin(c)

INLINE struct temp_arena _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 = thread_local_eval(&tl_scratch_ctx);
    struct arena *scratch_arena = &ctx->arenas[0];
    if (potential_conflict && scratch_arena->base == potential_conflict->base) {
        scratch_arena = &ctx->arenas[1];
    }
    struct temp_arena 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 context (`scratch_begin(arena)` should be called
 * instead). */
#define scratch_begin_no_conflict() \
    _scratch_begin_no_conflict(); \
    do { \
        struct arena *arena = NULL; \
        (UNUSED)arena; \
    } while (0)

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

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

INLINE void scratch_end(struct temp_arena scratch_temp)
{
#if RTC
    struct scratch_ctx *ctx = thread_local_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);
}

INLINE void scratch_end_and_decommit(struct temp_arena scratch_temp)
{
    scratch_end(scratch_temp);

    /* Disabled for now */
    // arena_decommit_unused_blocks(scratch_temp.arena);
}

#endif