/*
 * JIT compiler for RISC-V RV64I -> x86-64.
 *
 * Translates basic blocks of RISC-V instructions into native x86-64 machine
 * code.  The generated code operates on the guest CPU state structure directly
 * and falls back to the interpreter for ecalls, ebreak, and faults.
 *
 * Design follows QEMU's TCG approach: translate on first encounter, cache the
 * result, and execute the native code on subsequent visits.
 */
#ifndef JIT_H
#define JIT_H

#include "riscv.h"
#include "types.h"
#include <stdbool.h>

/* Forward declarations. */
struct cpu;

/* ---------------------------------------------------------------------------
 * Tuning constants
 * ------------------------------------------------------------------------- */

/* Size of the JIT code cache (16 MB). */
#define JIT_CODE_CACHE_SIZE (16 * 1024 * 1024)

/* Maximum number of x86-64 bytes emitted per guest instruction.
 * Conservative upper bound -- most instructions compile to <30 bytes. */
#define JIT_MAX_INSN_BYTES 128

/* Maximum guest instructions in a single translated block. */
#define JIT_MAX_BLOCK_INSNS 256

/* Number of hash-table buckets for block lookup (must be power of 2). */
#define JIT_BLOCK_HASH_SIZE (64 * 1024)

/* Maximum number of translated blocks. */
#define JIT_MAX_BLOCKS (256 * 1024)

/* Exit reasons returned by generated code. */
enum jit_exit {
    JIT_EXIT_BRANCH = 0, /* End of block -- branch/jump, update PC. */
    JIT_EXIT_ECALL  = 1, /* ECALL -- needs interpreter handling. */
    JIT_EXIT_EBREAK = 2, /* EBREAK -- needs interpreter handling. */
    JIT_EXIT_FAULT  = 3, /* Memory fault -- needs interpreter handling. */
    JIT_EXIT_CHAIN  = 4, /* Chain to next block (PC already set). */
    JIT_EXIT_RET    = 5, /* Program returned (jalr x0, ra, 0 with ra==0). */
};

/* A translated block of native code. */
struct jit_block {
    u32               guest_pc;     /* Guest PC this block starts at. */
    u32               guest_end_pc; /* Guest PC *after* last instruction. */
    u32               insn_count; /* Number of guest instructions translated.*/
    u8               *code;       /* Pointer into code cache. */
    u32               code_size;  /* Size of generated x86-64 code. */
    struct jit_block *hash_next;  /* Hash chain. */
};

/* Complete JIT state. */
struct jit_state {
    /* Executable code cache (mmap'd RWX). */
    u8 *code_cache;
    u32 code_cache_used;

    /* Block hash table (open-chaining). */
    struct jit_block *block_hash[JIT_BLOCK_HASH_SIZE];

    /* Block storage (static pool). */
    struct jit_block blocks[JIT_MAX_BLOCKS];
    u32              block_count;

    /* Statistics. */
    u64 blocks_compiled;
    u64 blocks_executed;
    u64 insns_executed;

    /* Whether JIT is available (mmap succeeded, x86-64 host). */
    bool available;
};

/* ---------------------------------------------------------------------------
 * API
 * ------------------------------------------------------------------------- */

/* Initialise the JIT.  Returns true on success.
 * On non-x86-64 hosts or if mmap fails, returns false and the emulator
 * should fall back to the interpreter. */
bool jit_init(struct jit_state *jit);

/* Release JIT resources. */
void jit_destroy(struct jit_state *jit);

/* Flush the entire translation cache (e.g. after self-modifying code). */
void jit_flush(struct jit_state *jit);

/* Compile (but don't look up) a block starting at `guest_pc`.
 * Returns NULL if the code cache is full or compilation fails. */
struct jit_block *jit_compile_block(
    struct jit_state *jit,
    u32               guest_pc,
    u8               *memory,
    u32               program_base,
    u32               program_bytes
);

/* Inline fast-path block lookup (hash table probe). */
static inline struct jit_block *jit_lookup_block(
    struct jit_state *jit, u32 guest_pc
) {
    u32               h = (guest_pc >> 2) & (JIT_BLOCK_HASH_SIZE - 1);
    struct jit_block *b = jit->block_hash[h];
    while (b) {
        if (b->guest_pc == guest_pc)
            return b;
        b = b->hash_next;
    }
    return NULL;
}

/* Look up or compile a block.  Calls the inline lookup first, falls
 * back to the compiler on a miss. */
static inline struct jit_block *jit_get_block(
    struct jit_state *jit,
    u32               guest_pc,
    u8               *memory,
    u32               program_base,
    u32               program_bytes
) {
    struct jit_block *b = jit_lookup_block(jit, guest_pc);
    if (b)
        return b;
    return jit_compile_block(
        jit, guest_pc, memory, program_base, program_bytes
    );
}

/* Signature of generated block entry points.
 * Args: pointer to cpu->regs[], pointer to memory[].
 * Returns: exit reason (enum jit_exit). */
typedef int (*jit_block_fn)(u64 *regs, u8 *memory, u32 *pc_out);

/* Execute a translated block.  Returns the exit reason. */
static inline int jit_exec_block(
    struct jit_block *block, u64 *regs, u8 *memory, u32 *pc_out
) {
    jit_block_fn fn = (jit_block_fn)(void *)block->code;
    return fn(regs, memory, pc_out);
}

#endif /* JIT_H */