#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 "gguf.h"
#include "weight-ctx.h"

#include <cstdio>
#include <cstdlib>
#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);

    // Verify every tensor fits inside the mapped file. Catches truncated
    // downloads early with a clear message instead of a segfault deep in
    // cuMemcpyHtoDAsync when the backend reads past the mmap.
    for (int64_t i = 0; i < n; i++) {
        const char *         tname = gguf_get_tensor_name(gf->gguf, i);
        struct ggml_tensor * t     = ggml_get_tensor(gf->meta, tname);
        size_t               toff  = gguf_get_tensor_offset(gf->gguf, i);
        size_t               tsize = ggml_nbytes(t);
        size_t               end   = gf->data_offset + toff + tsize;
        if (end > gf->file_size) {
            fprintf(stderr,
                    "[GGUF] FATAL: '%s' is truncated or corrupt.\n"
                    "       tensor '%s' needs bytes [%zu..%zu) but file is only %zu bytes.\n"
                    "       Re-download the file and verify its size or checksum.\n",
                    path, tname, gf->data_offset + toff, end, gf->file_size);
            gf_close(gf);
            return false;
        }
    }

    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. unique_ptr keeps .get() stable even
    // when wctx->staging grows on subsequent calls.
    size_t  n    = ggml_nelements(src);
    auto    buf  = std::make_unique<float[]>(n);
    float * data = buf.get();

    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++) {
            data[i] = ggml_bf16_to_fp32(*(const ggml_bf16_t *) &p[i]);
        }
    } else {
        ggml_fp16_to_fp32_row((const ggml_fp16_t *) raw, data, (int) n);
    }

    wctx->pending.push_back({ tensor, data, n * sizeof(float), 0 });
    wctx->staging.push_back(std::move(buf));
    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);
}