#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#if defined(_WIN32)
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#else
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#endif

#include "bio/bio_gguf.h"

#define BIO_GGUF_MAGIC 0x46554747u
#define BIO_GGUF_DEFAULT_ALIGNMENT 32u

typedef struct bio_reader {
    const uint8_t *base;
    size_t size;
    size_t off;
} bio_reader;

static uint64_t align_up_u64(uint64_t value, uint32_t alignment) {
    const uint64_t mask = (uint64_t)alignment - 1u;
    return (value + mask) & ~mask;
}

static int read_bytes(bio_reader *r, void *out, size_t n) {
    if (r == NULL || out == NULL || r->off > r->size || n > r->size - r->off) {
        return 0;
    }
    memcpy(out, r->base + r->off, n);
    r->off += n;
    return 1;
}

static int read_u32(bio_reader *r, uint32_t *out) {
    uint8_t b[4];
    if (!read_bytes(r, b, sizeof(b))) {
        return 0;
    }
    *out = (uint32_t)b[0] | ((uint32_t)b[1] << 8) | ((uint32_t)b[2] << 16) |
           ((uint32_t)b[3] << 24);
    return 1;
}

static int read_u64(bio_reader *r, uint64_t *out) {
    uint8_t b[8];
    if (!read_bytes(r, b, sizeof(b))) {
        return 0;
    }
    *out = (uint64_t)b[0] | ((uint64_t)b[1] << 8) | ((uint64_t)b[2] << 16) |
           ((uint64_t)b[3] << 24) | ((uint64_t)b[4] << 32) |
           ((uint64_t)b[5] << 40) | ((uint64_t)b[6] << 48) |
           ((uint64_t)b[7] << 56);
    return 1;
}

static int read_string_ref(bio_reader *r, const uint8_t **data_out,
                           uint64_t *len_out) {
    uint64_t len = 0;
    if (!read_u64(r, &len)) {
        return 0;
    }
    if (len > (uint64_t)(r->size - r->off)) {
        return 0;
    }
    *data_out = r->base + r->off;
    *len_out = len;
    r->off += (size_t)len;
    return 1;
}

static void copy_string_ref(char *dst, size_t cap, const uint8_t *src, uint64_t len) {
    const size_t n = len < (uint64_t)(cap - 1u) ? (size_t)len : cap - 1u;
    memcpy(dst, src, n);
    dst[n] = '\0';
}

static int key_equals(const uint8_t *key, uint64_t len, const char *literal) {
    const size_t n = strlen(literal);
    return len == (uint64_t)n && memcmp(key, literal, n) == 0;
}

static int skip_value(bio_reader *r, uint32_t type, uint32_t depth);

static int skip_array(bio_reader *r, uint32_t depth) {
    uint32_t elem_type = 0;
    uint64_t count = 0;
    if (depth > 4u || !read_u32(r, &elem_type) || !read_u64(r, &count)) {
        return 0;
    }
    for (uint64_t i = 0; i < count; ++i) {
        if (!skip_value(r, elem_type, depth + 1u)) {
            return 0;
        }
    }
    return 1;
}

static int skip_value(bio_reader *r, uint32_t type, uint32_t depth) {
    const uint8_t *s = NULL;
    uint64_t len = 0;
    switch (type) {
    case 0: /* uint8 */
    case 1: /* int8 */
    case 7: /* bool */
        r->off += 1u;
        return r->off <= r->size;
    case 2: /* uint16 */
    case 3: /* int16 */
        r->off += 2u;
        return r->off <= r->size;
    case 4: /* uint32 */
    case 5: /* int32 */
    case 6: /* float32 */
        r->off += 4u;
        return r->off <= r->size;
    case 8:  /* string */
        return read_string_ref(r, &s, &len);
    case 9:  /* array */
        return skip_array(r, depth);
    case 10: /* uint64 */
    case 11: /* int64 */
    case 12: /* float64 */
        r->off += 8u;
        return r->off <= r->size;
    default:
        return 0;
    }
}

static int read_metadata_u32(bio_reader *r, uint32_t type, uint32_t *out) {
    uint32_t v32 = 0;
    uint64_t v64 = 0;
    if (type == 4u) {
        if (!read_u32(r, &v32)) {
            return 0;
        }
        *out = v32;
        return 1;
    }
    if (type == 10u) {
        if (!read_u64(r, &v64) || v64 > 0xffffffffull) {
            return 0;
        }
        *out = (uint32_t)v64;
        return 1;
    }
    return 0;
}

static int read_token_array(bio_reader *r, bio_gguf_model *model, uint32_t depth) {
    uint32_t elem_type = 0;
    uint64_t count = 0;
    if (depth > 4u || !read_u32(r, &elem_type) || !read_u64(r, &count) ||
        elem_type != 8u || count > BIO_GGUF_MAX_TOKENS) {
        return 0;
    }
    model->tokenizer_token_count = (uint32_t)count;
    model->vocab_size = (uint32_t)count;
    for (uint64_t i = 0; i < count; ++i) {
        if (!read_string_ref(r, &model->tokenizer_tokens[i].data,
                             &model->tokenizer_tokens[i].len)) {
            return 0;
        }
    }
    return 1;
}

static uint64_t ggml_type_block_size(uint32_t type) {
    switch (type) {
    case BIO_GGUF_TYPE_Q4_0:
    case BIO_GGUF_TYPE_Q4_1:
    case BIO_GGUF_TYPE_Q5_0:
    case BIO_GGUF_TYPE_Q5_1:
    case BIO_GGUF_TYPE_Q8_0:
        return 32u;
    case BIO_GGUF_TYPE_Q2_K:
    case BIO_GGUF_TYPE_Q3_K:
    case BIO_GGUF_TYPE_Q4_K:
    case BIO_GGUF_TYPE_Q5_K:
    case BIO_GGUF_TYPE_Q6_K:
        return 256u;
    case BIO_GGUF_TYPE_IQ4_NL:
        return 32u;
    default:
        return 1u;
    }
}

static uint64_t ggml_type_bytes_per_block(uint32_t type) {
    switch (type) {
    case BIO_GGUF_TYPE_F32:
        return 4u;
    case BIO_GGUF_TYPE_F16:
        return 2u;
    case BIO_GGUF_TYPE_Q4_0:
        return 18u;
    case BIO_GGUF_TYPE_Q4_1:
        return 20u;
    case BIO_GGUF_TYPE_Q5_0:
        return 22u;
    case BIO_GGUF_TYPE_Q5_1:
        return 24u;
    case BIO_GGUF_TYPE_Q8_0:
        return 34u;
    case BIO_GGUF_TYPE_Q2_K:
        return 84u;
    case BIO_GGUF_TYPE_Q3_K:
        return 110u;
    case BIO_GGUF_TYPE_Q4_K:
        return 144u;
    case BIO_GGUF_TYPE_Q5_K:
        return 176u;
    case BIO_GGUF_TYPE_Q6_K:
        return 210u;
    case BIO_GGUF_TYPE_IQ4_NL:
        return 18u;
    default:
        return 0u;
    }
}

static uint64_t tensor_byte_size(const bio_gguf_tensor_span *tensor) {
    const uint64_t elements = bio_gguf_tensor_element_count(tensor);
    const uint64_t block = ggml_type_block_size(tensor->ggml_type);
    const uint64_t bytes = ggml_type_bytes_per_block(tensor->ggml_type);
    if (elements == 0u || block == 0u || bytes == 0u) {
        return 0u;
    }
    return ((elements + block - 1u) / block) * bytes;
}

static void copy_name(char *dst, size_t cap, const uint8_t *src, uint64_t len) {
    const size_t n = len < (uint64_t)(cap - 1u) ? (size_t)len : cap - 1u;
    memcpy(dst, src, n);
    dst[n] = '\0';
}

static bio_v1_status parse_gguf_mapping(bio_gguf_model *model) {
    bio_reader r;
    uint32_t magic = 0;

    memset(model->tensors, 0, sizeof(model->tensors));
    model->alignment = BIO_GGUF_DEFAULT_ALIGNMENT;
    model->registered_tensor_count = 0;

    r.base = (const uint8_t *)model->mapping;
    r.size = model->mapping_size;
    r.off = 0;

    if (!read_u32(&r, &magic) || magic != BIO_GGUF_MAGIC ||
        !read_u32(&r, &model->version) ||
        !read_u64(&r, &model->tensor_count) ||
        !read_u64(&r, &model->metadata_count)) {
        return BIO_ERR_PARSE;
    }
    if (model->version < 2u || model->version > 3u ||
        model->tensor_count > BIO_GGUF_MAX_TENSORS) {
        return BIO_ERR_UNSUPPORTED;
    }

    for (uint64_t i = 0; i < model->metadata_count; ++i) {
        const uint8_t *key = NULL;
        uint64_t key_len = 0;
        uint32_t type = 0;
        if (!read_string_ref(&r, &key, &key_len) || !read_u32(&r, &type)) {
            return BIO_ERR_PARSE;
        }
        if (key_equals(key, key_len, "general.alignment") && type == 4u) {
            if (!read_u32(&r, &model->alignment)) {
                return BIO_ERR_PARSE;
            }
            if (model->alignment == 0u || (model->alignment & (model->alignment - 1u)) != 0u) {
                return BIO_ERR_UNSUPPORTED;
            }
            continue;
        }
        if (key_equals(key, key_len, "general.architecture") && type == 8u) {
            const uint8_t *value = NULL;
            uint64_t value_len = 0;
            if (!read_string_ref(&r, &value, &value_len)) {
                return BIO_ERR_PARSE;
            }
            copy_string_ref(model->architecture, sizeof(model->architecture), value, value_len);
            continue;
        }
        if ((key_equals(key, key_len, "phi3.context_length") ||
             key_equals(key, key_len, "llama.context_length")) &&
            read_metadata_u32(&r, type, &model->context_length)) {
            continue;
        }
        if ((key_equals(key, key_len, "phi3.embedding_length") ||
             key_equals(key, key_len, "llama.embedding_length")) &&
            read_metadata_u32(&r, type, &model->embedding_length)) {
            continue;
        }
        if ((key_equals(key, key_len, "phi3.block_count") ||
             key_equals(key, key_len, "llama.block_count")) &&
            read_metadata_u32(&r, type, &model->block_count)) {
            continue;
        }
        if ((key_equals(key, key_len, "phi3.feed_forward_length") ||
             key_equals(key, key_len, "llama.feed_forward_length")) &&
            read_metadata_u32(&r, type, &model->feed_forward_length)) {
            continue;
        }
        if ((key_equals(key, key_len, "phi3.attention.head_count") ||
             key_equals(key, key_len, "llama.attention.head_count")) &&
            read_metadata_u32(&r, type, &model->attention_head_count)) {
            continue;
        }
        if ((key_equals(key, key_len, "phi3.attention.head_count_kv") ||
             key_equals(key, key_len, "llama.attention.head_count_kv")) &&
            read_metadata_u32(&r, type, &model->attention_head_count_kv)) {
            continue;
        }
        if (key_equals(key, key_len, "tokenizer.ggml.tokens") && type == 9u) {
            if (!read_token_array(&r, model, 0u)) {
                return BIO_ERR_PARSE;
            }
            continue;
        }
        if (!skip_value(&r, type, 0u)) {
            return BIO_ERR_PARSE;
        }
    }

    for (uint64_t i = 0; i < model->tensor_count; ++i) {
        bio_gguf_tensor_span *tensor = &model->tensors[i];
        const uint8_t *name = NULL;
        uint64_t name_len = 0;
        uint64_t tensor_offset = 0;
        if (!read_string_ref(&r, &name, &name_len) ||
            !read_u32(&r, &tensor->dims_count)) {
            return BIO_ERR_PARSE;
        }
        if (tensor->dims_count == 0u || tensor->dims_count > BIO_GGUF_MAX_DIMS) {
            return BIO_ERR_UNSUPPORTED;
        }
        copy_name(tensor->name, sizeof(tensor->name), name, name_len);
        for (uint32_t d = 0; d < tensor->dims_count; ++d) {
            if (!read_u64(&r, &tensor->dims[d])) {
                return BIO_ERR_PARSE;
            }
        }
        if (!read_u32(&r, &tensor->ggml_type) || !read_u64(&r, &tensor_offset)) {
            return BIO_ERR_PARSE;
        }
        if (ggml_type_bytes_per_block(tensor->ggml_type) == 0u) {
            return BIO_ERR_UNSUPPORTED;
        }
        tensor->data_offset = tensor_offset;
        tensor->int4_compatible = (uint8_t)bio_gguf_type_is_int4(tensor->ggml_type);
        model->registered_tensor_count += 1u;
    }

    {
        const uint64_t data_base = align_up_u64((uint64_t)r.off, model->alignment);
        for (uint64_t i = 0; i < model->tensor_count; ++i) {
            bio_gguf_tensor_span *tensor = &model->tensors[i];
            const uint64_t size = tensor_byte_size(tensor);
            const uint64_t absolute = data_base + tensor->data_offset;
            if (size == 0u || absolute > (uint64_t)model->mapping_size ||
                size > (uint64_t)model->mapping_size - absolute) {
                return BIO_ERR_PARSE;
            }
            tensor->file_offset = absolute;
            tensor->byte_size = size;
            tensor->data = (const uint8_t *)model->mapping + absolute;
        }
    }

    return BIO_OK;
}

int bio_gguf_type_is_int4(uint32_t ggml_type) {
    return ggml_type == BIO_GGUF_TYPE_Q4_0 || ggml_type == BIO_GGUF_TYPE_Q4_1 ||
           ggml_type == BIO_GGUF_TYPE_Q4_K || ggml_type == BIO_GGUF_TYPE_IQ4_NL;
}

uint64_t bio_gguf_tensor_element_count(const bio_gguf_tensor_span *tensor) {
    uint64_t n = 1u;
    if (tensor == NULL || tensor->dims_count == 0u) {
        return 0u;
    }
    for (uint32_t i = 0; i < tensor->dims_count; ++i) {
        if (tensor->dims[i] == 0u || n > UINT64_MAX / tensor->dims[i]) {
            return 0u;
        }
        n *= tensor->dims[i];
    }
    return n;
}

const bio_gguf_tensor_span *bio_gguf_find_tensor(const bio_gguf_model *model,
                                                 const char *name) {
    if (model == NULL || name == NULL) {
        return NULL;
    }
    for (uint32_t i = 0; i < model->registered_tensor_count; ++i) {
        if (strcmp(model->tensors[i].name, name) == 0) {
            return &model->tensors[i];
        }
    }
    return NULL;
}

uint32_t bio_tokenize_bytes(const char *text, uint32_t *tokens,
                            uint32_t max_tokens) {
    uint32_t n = 0;
    if (text == NULL || tokens == NULL || max_tokens == 0u) {
        return 0u;
    }
    while (*text != '\0' && n < max_tokens) {
        tokens[n++] = (uint8_t)*text++;
    }
    return n;
}

uint32_t bio_detokenize_bytes(const uint32_t *tokens, uint32_t token_count,
                              char *out, uint32_t out_cap) {
    uint32_t n = 0;
    if (tokens == NULL || out == NULL || out_cap == 0u) {
        return 0u;
    }
    for (; n + 1u < out_cap && n < token_count; ++n) {
        const uint32_t t = tokens[n] & 0xffu;
        out[n] = (char)(t >= 32u && t < 127u ? t : '.');
    }
    out[n] = '\0';
    return n;
}

bio_v1_status bio_gguf_open(const char *path, bio_gguf_model *model) {
    bio_v1_status rc = BIO_OK;
    if (path == NULL || model == NULL || *path == '\0') {
        return BIO_ERR_ARGS;
    }
    memset(model, 0, sizeof(*model));
#if defined(_WIN32)
    model->file_handle = CreateFileA(path, GENERIC_READ, FILE_SHARE_READ, NULL,
                                     OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (model->file_handle == INVALID_HANDLE_VALUE) {
        model->file_handle = NULL;
        return BIO_ERR_IO;
    }
    {
        LARGE_INTEGER size;
        if (!GetFileSizeEx((HANDLE)model->file_handle, &size) || size.QuadPart <= 0 ||
            (uint64_t)size.QuadPart > (uint64_t)SIZE_MAX) {
            bio_gguf_close(model);
            return BIO_ERR_IO;
        }
        model->mapping_size = (size_t)size.QuadPart;
    }
    model->mapping_handle =
        CreateFileMappingA((HANDLE)model->file_handle, NULL, PAGE_READONLY, 0, 0, NULL);
    if (model->mapping_handle == NULL) {
        bio_gguf_close(model);
        return BIO_ERR_IO;
    }
    model->mapping = MapViewOfFile((HANDLE)model->mapping_handle, FILE_MAP_READ, 0, 0, 0);
    if (model->mapping == NULL) {
        bio_gguf_close(model);
        return BIO_ERR_IO;
    }
#else
    model->fd = -1;
    model->fd = open(path, O_RDONLY);
    if (model->fd < 0) {
        return BIO_ERR_IO;
    }
    {
        struct stat st;
        if (fstat(model->fd, &st) != 0 || st.st_size <= 0 ||
            (uint64_t)st.st_size > (uint64_t)SIZE_MAX) {
            bio_gguf_close(model);
            return BIO_ERR_IO;
        }
        model->mapping_size = (size_t)st.st_size;
    }
    model->mapping = mmap(NULL, model->mapping_size, PROT_READ, MAP_PRIVATE, model->fd, 0);
    if (model->mapping == MAP_FAILED) {
        model->mapping = NULL;
        bio_gguf_close(model);
        return BIO_ERR_IO;
    }
#endif
    rc = parse_gguf_mapping(model);
    if (rc != BIO_OK) {
        bio_gguf_close(model);
    }
    return rc;
}

void bio_gguf_close(bio_gguf_model *model) {
    if (model == NULL) {
        return;
    }
#if defined(_WIN32)
    if (model->mapping != NULL) {
        UnmapViewOfFile(model->mapping);
    }
    if (model->mapping_handle != NULL) {
        CloseHandle((HANDLE)model->mapping_handle);
    }
    if (model->file_handle != NULL) {
        CloseHandle((HANDLE)model->file_handle);
    }
#else
    if (model->mapping != NULL) {
        munmap(model->mapping, model->mapping_size);
    }
    if (model->fd >= 0) {
        close(model->fd);
    }
#endif
    memset(model, 0, sizeof(*model));
#if !defined(_WIN32)
    model->fd = -1;
#endif
}
