#pragma once

#include <stddef.h>
#include <stdint.h>

#ifdef __cplusplus
extern "C" {
#endif

/* ── File format constants ───────────────────────────────────────
 *
 * On-disk layout (little-endian throughout):
 *
 *   [ magic    : char[4]   ]  "MEMB"
 *   [ version  : uint32    ]  MEMBA_FILE_VERSION
 *   [ model_id : char[64]  ]  hex-encoded SHA-256 of first 1 KiB of GGUF
 *   [ n_ctx    : uint32    ]  llama_n_ctx() at save time
 *   [ llama_ver: uint32    ]  reserved (0 for now)
 *   [ data_size: uint64    ]  byte length of the opaque state blob
 *   [ data     : uint8[]   ]  llama_state_get_data() blob
 *   [ crc32    : uint32    ]  CRC-32 of data[] only
 */
#define MEMBA_FILE_MAGIC    "MEMB"
#define MEMBA_FILE_VERSION  1u
#define MEMBA_MODEL_ID_LEN  64

/* ── Error codes ─────────────────────────────────────────────────*/
#define MEMBA_OK            0
#define MEMBA_ERR_IO       -1   /* file open / read / write failed  */
#define MEMBA_ERR_MAGIC    -2   /* bad magic bytes                  */
#define MEMBA_ERR_VERSION  -3   /* unsupported file version         */
#define MEMBA_ERR_MODEL_ID -4   /* model identity mismatch          */
#define MEMBA_ERR_CRC      -5   /* CRC-32 checksum mismatch         */
#define MEMBA_ERR_ALLOC    -6   /* memory allocation failed         */
#define MEMBA_ERR_CTX      -7   /* null or invalid llama_context    */

/* Opaque handle — one per llama_context you want to checkpoint. */
typedef struct memba_state memba_state_t;

struct llama_context; /* forward declaration */

/**
 * Create a handle that wraps @p ctx.
 *
 * @param ctx        Active llama_context. Must remain alive for the handle's
 *                   entire lifetime.
 * @param model_path Path to the GGUF file. Used to compute the model identity
 *                   fingerprint embedded in every state file. Pass NULL to
 *                   skip identity checking (strongly discouraged).
 * @return           New handle, or NULL on allocation failure.
 */
memba_state_t* memba_state_new(struct llama_context* ctx, const char* model_path);

/** Free a handle created by memba_state_new(). Does NOT free ctx. */
void memba_state_free(memba_state_t* state);

/**
 * Serialise the current SSM hidden state to @p path.
 *
 * Thread-safety: the caller must ensure no concurrent llama_decode() calls
 * on @p ctx while this function executes.
 *
 * @return MEMBA_OK on success, negative error code on failure.
 */
int memba_state_save(memba_state_t* state, const char* path);

/**
 * Restore state from @p path into the wrapped llama_context.
 *
 * Validates magic, version, model_id, and CRC-32 before applying.
 *
 * @return MEMBA_OK on success, negative error code on failure.
 */
int memba_state_load(memba_state_t* state, const char* path);

/**
 * Return the serialised byte size of the current state (useful for logging).
 * Returns 0 if the handle is NULL or ctx is invalid.
 */
size_t memba_state_get_size(memba_state_t* state);

/** Human-readable description of an error code. Never returns NULL. */
const char* memba_error_string(int err);

#ifdef __cplusplus
}
#endif