#pragma once
// gguf_weights.h: load model weights from GGUF files
//
// GGUF weight loader for all model components (LM, DiT, CondEncoder, TextEncoder, Detokenizer, VAE).
// All components use GGUF bf16 files generated by convert.py.
//
// Usage:
//   GGUFModel gf;
//   if (!gf_load(&gf, "model.gguf")) { error; }
//   WeightCtx wctx;
//   wctx_init(&wctx, n_tensors);
//   ggml_tensor * w = gf_load_tensor(&wctx, gf, "layer.0.weight");
//   wctx_alloc(&wctx, backend);
//   gf_close(&gf);   // safe after wctx_alloc copied data to GPU

#include "weight_ctx.h"
#include "gguf.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>

#ifdef _WIN32
#include <windows.h>
#else
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#endif

struct GGUFModel {
    struct gguf_context * gguf;     // parsed header (KV + tensor metadata)
    struct ggml_context * meta;     // tensor descriptors (no data)
    uint8_t * mapping;              // mmapped file
    size_t file_size;
    size_t data_offset;             // gguf_get_data_offset(gguf)
#ifdef _WIN32
    HANDLE fh;
    HANDLE mh;
#else
    int fd;
#endif
};

static void gf_close(GGUFModel * gf) {
    if (gf->gguf) gguf_free(gf->gguf);
    if (gf->meta) ggml_free(gf->meta);
#ifdef _WIN32
    if (gf->mapping) UnmapViewOfFile(gf->mapping);
    if (gf->mh) CloseHandle(gf->mh);
    if (gf->fh && gf->fh != INVALID_HANDLE_VALUE) CloseHandle(gf->fh);
#else
    if (gf->mapping) munmap(gf->mapping, gf->file_size);
    if (gf->fd >= 0) close(gf->fd);
#endif
    *gf = {};
}

static bool gf_load(GGUFModel * gf, const char * path) {
    *gf = {};

    // mmap the file
#ifdef _WIN32
    gf->fh = CreateFileA(path, GENERIC_READ, FILE_SHARE_READ, NULL,
                          OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (gf->fh == INVALID_HANDLE_VALUE) {
        fprintf(stderr, "[GGUF] cannot open %s\n", path);
        return false;
    }
    LARGE_INTEGER li;
    GetFileSizeEx(gf->fh, &li);
    gf->file_size = (size_t)li.QuadPart;
    gf->mh = CreateFileMappingA(gf->fh, NULL, PAGE_READONLY, 0, 0, NULL);
    if (!gf->mh) {
        CloseHandle(gf->fh);
        fprintf(stderr, "[GGUF] CreateFileMapping failed %s\n", path);
        return false;
    }
    gf->mapping = (uint8_t *)MapViewOfFile(gf->mh, FILE_MAP_READ, 0, 0, 0);
    if (!gf->mapping) {
        CloseHandle(gf->mh);
        CloseHandle(gf->fh);
        fprintf(stderr, "[GGUF] MapViewOfFile failed %s\n", path);
        return false;
    }
#else
    gf->fd = open(path, O_RDONLY);
    if (gf->fd < 0) {
        fprintf(stderr, "[GGUF] cannot open %s\n", path);
        return false;
    }
    struct stat sb;
    fstat(gf->fd, &sb);
    gf->file_size = (size_t)sb.st_size;
    gf->mapping = (uint8_t *)mmap(NULL, gf->file_size, PROT_READ, MAP_PRIVATE, gf->fd, 0);
    if (gf->mapping == MAP_FAILED) {
        close(gf->fd);
        gf->mapping = NULL;
        fprintf(stderr, "[GGUF] mmap failed %s\n", path);
        return false;
    }
#endif

    // Parse GGUF header, create tensor metadata context
    struct ggml_context * meta = NULL;
    struct gguf_init_params params = { /*no_alloc=*/ true, /*ctx=*/ &meta };
    gf->gguf = gguf_init_from_file(path, params);
    if (!gf->gguf) {
        fprintf(stderr, "[GGUF] failed to parse %s\n", path);
        gf_close(gf);
        return false;
    }
    gf->meta = meta;
    gf->data_offset = gguf_get_data_offset(gf->gguf);

    int64_t n = gguf_get_n_tensors(gf->gguf);
    fprintf(stderr, "[GGUF] %s: %lld tensors, data at offset %zu\n",
            path, (long long)n, gf->data_offset);
    return true;
}

// Load a tensor from GGUF into the weight context.
// Returns ggml_tensor (not yet backed by memory; call wctx_alloc after all loads).
// Tensor shapes are already in ggml order (ne[0]=innermost).
static struct ggml_tensor * gf_load_tensor(
        WeightCtx * wctx,
        const GGUFModel & gf,
        const std::string & name,
        const int64_t * shape_override = nullptr,
        int n_dims_override = 0) {

    int64_t idx = gguf_find_tensor(gf.gguf, name.c_str());
    if (idx < 0) {
        fprintf(stderr, "[GGUF] FATAL: tensor '%s' not found\n", name.c_str());
        exit(1);
    }

    // Get metadata from the context populated by gguf_init_from_file
    struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str());
    if (!src) {
        fprintf(stderr, "[GGUF] FATAL: tensor '%s' not in meta context\n", name.c_str());
        exit(1);
    }

    int n_dims;
    int64_t ne[4] = {1, 1, 1, 1};

    if (shape_override && n_dims_override > 0) {
        n_dims = n_dims_override;
        for (int i = 0; i < n_dims; i++) ne[i] = shape_override[i];
    } else {
        n_dims = ggml_n_dims(src);
        for (int i = 0; i < n_dims; i++) ne[i] = src->ne[i];
    }

    struct ggml_tensor * tensor = ggml_new_tensor(wctx->ctx, src->type, n_dims, ne);
    ggml_set_name(tensor, name.c_str());

    size_t offset = gguf_get_tensor_offset(gf.gguf, idx);
    const void * data = gf.mapping + gf.data_offset + offset;
    size_t nbytes = ggml_nbytes(src);

    wctx->pending.push_back({tensor, data, nbytes, 0});
    return tensor;
}

// Try to load, returns nullptr if not found (no exit)
static struct ggml_tensor * gf_try_load_tensor(
        WeightCtx * wctx,
        const GGUFModel & gf,
        const std::string & name) {
    int64_t idx = gguf_find_tensor(gf.gguf, name.c_str());
    if (idx < 0) return nullptr;
    return gf_load_tensor(wctx, gf, name);
}

// Load tensor, converting to F32 at load time (eliminates runtime cast nodes).
// Best for small tensors: norms [H], QK-norms [D], scale_shift_table [H,6], biases.
static struct ggml_tensor * gf_load_tensor_f32(
        WeightCtx * wctx,
        const GGUFModel & gf,
        const std::string & name) {
    int64_t idx = gguf_find_tensor(gf.gguf, name.c_str());
    if (idx < 0) {
        fprintf(stderr, "[GGUF] FATAL: tensor '%s' not found\n", name.c_str());
        exit(1);
    }
    struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str());
    int n_dims = ggml_n_dims(src);
    int64_t ne[4] = {1, 1, 1, 1};
    for (int i = 0; i < n_dims; i++) ne[i] = src->ne[i];

    // If already F32, just load normally
    if (src->type == GGML_TYPE_F32) {
        return gf_load_tensor(wctx, gf, name);
    }

    // Bail early on unsupported types (before creating tensor in ctx)
    if (src->type != GGML_TYPE_BF16 && src->type != GGML_TYPE_F16) {
        fprintf(stderr, "[GGUF] WARNING: gf_load_tensor_f32 unsupported type %d for '%s', loading as-is\n",
                src->type, name.c_str());
        return gf_load_tensor(wctx, gf, name);
    }

    // Create F32 tensor
    struct ggml_tensor * tensor = ggml_new_tensor(wctx->ctx, GGML_TYPE_F32, n_dims, ne);
    ggml_set_name(tensor, name.c_str());

    // Convert data into staging buffer
    size_t n = ggml_nelements(src);
    wctx->staging.emplace_back(n);
    std::vector<float> & buf = wctx->staging.back();

    size_t offset = gguf_get_tensor_offset(gf.gguf, idx);
    const void * raw = gf.mapping + gf.data_offset + offset;

    if (src->type == GGML_TYPE_BF16) {
        const uint16_t * p = (const uint16_t *)raw;
        for (size_t i = 0; i < n; i++)
            buf[i] = ggml_bf16_to_fp32(*(const ggml_bf16_t *)&p[i]);
    } else {
        ggml_fp16_to_fp32_row((const ggml_fp16_t *)raw, buf.data(), (int)n);
    }

    wctx->pending.push_back({tensor, buf.data(), n * sizeof(float), 0});
    return tensor;
}

// Get raw pointer to tensor data in the mmapped file.
// Useful for CPU-side operations (e.g. bf16 embed lookup for lyrics).
// Returns NULL if not found.
static const void * gf_get_data(const GGUFModel & gf, const char * name) {
    int64_t idx = gguf_find_tensor(gf.gguf, name);
    if (idx < 0) return NULL;
    size_t offset = gguf_get_tensor_offset(gf.gguf, idx);
    return gf.mapping + gf.data_offset + offset;
}

// Fuse Q, K, V projection weights into a single tensor [ne0, q_ne1 + k_ne1 + v_ne1].
// Works for any quantized type since quantization is per-row (along ne[0]).
// The fused tensor data is q rows || k rows || v rows (contiguous).
static struct ggml_tensor * gf_load_qkv_fused(
        WeightCtx * wctx,
        const GGUFModel & gf,
        const std::string & q_name,
        const std::string & k_name,
        const std::string & v_name) {
    struct ggml_tensor * q_src = ggml_get_tensor(gf.meta, q_name.c_str());
    struct ggml_tensor * k_src = ggml_get_tensor(gf.meta, k_name.c_str());
    struct ggml_tensor * v_src = ggml_get_tensor(gf.meta, v_name.c_str());
    if (!q_src || !k_src || !v_src) {
        fprintf(stderr, "[GGUF] FATAL: QKV tensor not found: %s / %s / %s\n",
                q_name.c_str(), k_name.c_str(), v_name.c_str());
        exit(1);
    }
    // All must share ne[0] (input dim) and type - otherwise can't fuse
    GGML_ASSERT(q_src->ne[0] == k_src->ne[0] && k_src->ne[0] == v_src->ne[0]);
    if (q_src->type != k_src->type || k_src->type != v_src->type) {
        return NULL;  // caller should fall back to separate loads
    }

    int64_t ne0 = q_src->ne[0];
    int64_t fused_ne1 = q_src->ne[1] + k_src->ne[1] + v_src->ne[1];
    int64_t ne[2] = { ne0, fused_ne1 };
    struct ggml_tensor * fused = ggml_new_tensor(wctx->ctx, q_src->type, 2, ne);

    size_t row_size = ggml_row_size(q_src->type, ne0);
    size_t q_bytes = q_src->ne[1] * row_size;
    size_t k_bytes = k_src->ne[1] * row_size;
    size_t v_bytes = v_src->ne[1] * row_size;

    auto get_data = [&](const std::string & name) -> const void * {
        int64_t idx = gguf_find_tensor(gf.gguf, name.c_str());
        size_t off = gguf_get_tensor_offset(gf.gguf, idx);
        return gf.mapping + gf.data_offset + off;
    };

    wctx->pending.push_back({fused, get_data(q_name), q_bytes, 0});
    wctx->pending.push_back({fused, get_data(k_name), k_bytes, q_bytes});
    wctx->pending.push_back({fused, get_data(v_name), v_bytes, q_bytes + k_bytes});
    return fused;
}

// Fuse two projection weights [ne0, a_ne1 + b_ne1] when types match.
// Returns NULL if types differ.
static struct ggml_tensor * gf_load_pair_fused(
        WeightCtx * wctx,
        const GGUFModel & gf,
        const std::string & a_name,
        const std::string & b_name) {
    struct ggml_tensor * a_src = ggml_get_tensor(gf.meta, a_name.c_str());
    struct ggml_tensor * b_src = ggml_get_tensor(gf.meta, b_name.c_str());
    if (!a_src || !b_src) return NULL;
    if (a_src->ne[0] != b_src->ne[0] || a_src->type != b_src->type) return NULL;

    int64_t ne0 = a_src->ne[0];
    int64_t ne[2] = { ne0, a_src->ne[1] + b_src->ne[1] };
    struct ggml_tensor * fused = ggml_new_tensor(wctx->ctx, a_src->type, 2, ne);

    size_t row_size = ggml_row_size(a_src->type, ne0);
    size_t a_bytes = a_src->ne[1] * row_size;
    size_t b_bytes = b_src->ne[1] * row_size;

    auto get_data = [&](const std::string & name) -> const void * {
        int64_t idx = gguf_find_tensor(gf.gguf, name.c_str());
        size_t off = gguf_get_tensor_offset(gf.gguf, idx);
        return gf.mapping + gf.data_offset + off;
    };

    wctx->pending.push_back({fused, get_data(a_name), a_bytes, 0});
    wctx->pending.push_back({fused, get_data(b_name), b_bytes, a_bytes});
    return fused;
}

// Read a uint32 KV value (returns 0 if not found)
static uint32_t gf_get_u32(const GGUFModel & gf, const char * key) {
    int64_t idx = gguf_find_key(gf.gguf, key);
    if (idx < 0) return 0;
    return gguf_get_val_u32(gf.gguf, idx);
}

// Read a float32 KV value (returns 0 if not found)
static float gf_get_f32(const GGUFModel & gf, const char * key) {
    int64_t idx = gguf_find_key(gf.gguf, key);
    if (idx < 0) return 0.0f;
    return gguf_get_val_f32(gf.gguf, idx);
}

// Read a string KV value (returns "" if not found)
static const char * gf_get_str(const GGUFModel & gf, const char * key) {
    int64_t idx = gguf_find_key(gf.gguf, key);
    if (idx < 0) return "";
    return gguf_get_val_str(gf.gguf, idx);
}

// Read a bool KV value (returns false if not found)
static bool gf_get_bool(const GGUFModel & gf, const char * key) {
    int64_t idx = gguf_find_key(gf.gguf, key);
    if (idx < 0) return false;
    return gguf_get_val_bool(gf.gguf, idx);
}