/* Native GGUF mmap and tensor span contracts. */
#ifndef BIO_GGUF_H_
#define BIO_GGUF_H_

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

#include "bio/bio.h"

#ifdef __cplusplus
extern "C" {
#endif

#define BIO_GGUF_MAX_TENSORS 4096u
#define BIO_GGUF_MAX_NAME_BYTES 128u
#define BIO_GGUF_MAX_DIMS 8u
#define BIO_GGUF_MAX_TOKENS 131072u
#define BIO_GGUF_MAX_ARCH_BYTES 32u

typedef struct bio_gguf_str_ref {
    const uint8_t *data;
    uint64_t len;
} bio_gguf_str_ref;

typedef enum bio_gguf_tensor_type {
    BIO_GGUF_TYPE_F32 = 0,
    BIO_GGUF_TYPE_F16 = 1,
    BIO_GGUF_TYPE_Q4_0 = 2,
    BIO_GGUF_TYPE_Q4_1 = 3,
    BIO_GGUF_TYPE_Q5_0 = 6,
    BIO_GGUF_TYPE_Q5_1 = 7,
    BIO_GGUF_TYPE_Q8_0 = 8,
    BIO_GGUF_TYPE_Q2_K = 10,
    BIO_GGUF_TYPE_Q3_K = 11,
    BIO_GGUF_TYPE_Q4_K = 12,
    BIO_GGUF_TYPE_Q5_K = 13,
    BIO_GGUF_TYPE_Q6_K = 14,
    BIO_GGUF_TYPE_IQ4_NL = 20
} bio_gguf_tensor_type;

typedef struct bio_gguf_tensor_span {
    char name[BIO_GGUF_MAX_NAME_BYTES];
    uint32_t dims_count;
    uint64_t dims[BIO_GGUF_MAX_DIMS];
    uint32_t ggml_type;
    uint64_t file_offset;
    uint64_t data_offset;
    uint64_t byte_size;
    const uint8_t *data;
    uint8_t int4_compatible;
} bio_gguf_tensor_span;

typedef struct bio_gguf_model {
    void *mapping;
    size_t mapping_size;
    uint32_t version;
    uint64_t tensor_count;
    uint64_t metadata_count;
    uint32_t alignment;
    uint32_t registered_tensor_count;
    char architecture[BIO_GGUF_MAX_ARCH_BYTES];
    uint32_t context_length;
    uint32_t embedding_length;
    uint32_t block_count;
    uint32_t feed_forward_length;
    uint32_t attention_head_count;
    uint32_t attention_head_count_kv;
    uint32_t vocab_size;
    uint32_t tokenizer_token_count;
    bio_gguf_str_ref tokenizer_tokens[BIO_GGUF_MAX_TOKENS];
    bio_gguf_tensor_span tensors[BIO_GGUF_MAX_TENSORS];
#if defined(_WIN32)
    void *file_handle;
    void *mapping_handle;
#else
    int fd;
#endif
} bio_gguf_model;

bio_v1_status bio_gguf_open(const char *path, bio_gguf_model *model);
void bio_gguf_close(bio_gguf_model *model);
int bio_gguf_type_is_int4(uint32_t ggml_type);
uint64_t bio_gguf_tensor_element_count(const bio_gguf_tensor_span *tensor);
const bio_gguf_tensor_span *bio_gguf_find_tensor(const bio_gguf_model *model,
                                                 const char *name);
uint32_t bio_tokenize_bytes(const char *text, uint32_t *tokens,
                            uint32_t max_tokens);
uint32_t bio_detokenize_bytes(const uint32_t *tokens, uint32_t token_count,
                              char *out, uint32_t out_cap);

#ifdef __cplusplus
}
#endif

#endif /* BIO_GGUF_H_ */
