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

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

#ifdef _WIN32
#    include "utf8.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
    std::wstring wpath = utf8_to_wide(path);
    gf->fh =
        CreateFileW(wpath.c_str(), 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        = CreateFileMappingW(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) {
        qt_throw("[GGUF] tensor '%s' not found", name.c_str());
    }

    // 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) {
        qt_throw("[GGUF] tensor '%s' not in meta context", name.c_str());
    }

    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) {
        qt_throw("[GGUF] tensor '%s' not found (f32 load)", name.c_str());
    }
    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;
}

// Load a Conv1d / Conv1dDW kernel weight, forcing F16 storage on the
// backend regardless of the source GGUF dtype.
//
// TODO upstream GGML: ggml_conv_1d and ggml_conv_1d_dw in
// ggml/src/ggml.c hardcode dst_type = GGML_TYPE_F16 in their internal
// ggml_im2col call (currently ggml.c lines around 4508 and 4542).
// ggml_conv_2d at the equivalent site uses the adaptive pattern
// dst_type = a->type (around line 4595). Two backend bugs follow.
//
//   1) CPU im2col dispatches on dst->type. The im2col_f16 path
//      asserts src0->type == GGML_TYPE_F16, so a F32 or BF16 conv
//      kernel crashes on CPU.
//   2) ggml_conv_1d lowers to ggml_mul_mat(im2col, reshape(weight))
//      where the weight ends up as src1, not src0. The Vulkan fast
//      path ggml_vk_get_dequantize_mul_mat_vec asserts b_type in
//      {F32, F16, Q8_1}, so a BF16 conv kernel crashes on Vulkan
//      even when im2col itself succeeds.
//
// Vulkan does ship a pipeline_im2col_f32 and CUDA handles both F32
// and F16 cleanly, so the fix upstream is to align ggml_conv_1d and
// ggml_conv_1d_dw on ggml_conv_2d's adaptive a->type pattern. Until
// that lands, the only safe assumption across CPU, CUDA, and Vulkan
// is F16 kernels everywhere, so we mirror ggml_conv_1d's hardcoded
// choice here and load every conv kernel as F16 regardless of source
// dtype. F16 source is a direct passthrough, F32 and BF16 widen
// through a F32 staging buffer.
static struct ggml_tensor * gf_load_conv(WeightCtx * wctx, const GGUFModel & gf, const std::string & name) {
    int64_t idx = gguf_find_tensor(gf.gguf, name.c_str());
    if (idx < 0) {
        qt_throw("[GGUF] tensor '%s' not found (conv load)", name.c_str());
    }
    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];
    }

    // F16 source: direct passthrough, no conversion.
    if (src->type == GGML_TYPE_F16) {
        return gf_load_tensor(wctx, gf, name);
    }
    if (src->type != GGML_TYPE_F32 && src->type != GGML_TYPE_BF16) {
        qt_throw("[GGUF] gf_load_conv unsupported source type %s for '%s'", ggml_type_name(src->type), name.c_str());
    }

    // Allocate F16 backend tensor in the WeightCtx graph.
    struct ggml_tensor * tensor = ggml_new_tensor(wctx->ctx, GGML_TYPE_F16, n_dims, ne);
    ggml_set_name(tensor, name.c_str());

    size_t       n       = (size_t) ggml_nelements(src);
    size_t       raw_off = gguf_get_tensor_offset(gf.gguf, idx);
    const void * raw     = gf.mapping + gf.data_offset + raw_off;

    // The staging vector owns float[] buffers to keep memory alive
    // until wctx_alloc copies it to the backend. n F16 elements
    // occupy n * 2 bytes, which fits in (n + 1) / 2 floats. The
    // pending entry references the same buffer reinterpreted as
    // ggml_fp16_t and carries the exact F16 byte count.
    size_t        n_floats = (n + 1) / 2;
    auto          buf      = std::make_unique<float[]>(n_floats);
    ggml_fp16_t * data     = (ggml_fp16_t *) buf.get();

    if (src->type == GGML_TYPE_F32) {
        ggml_fp32_to_fp16_row((const float *) raw, data, (int) n);
    } else {
        // BF16 source: widen to F32 first, then narrow to F16 in
        // one pass to preserve mantissa bits the BF16-to-F16 direct
        // cast would otherwise leave undefined.
        std::vector<float> f32(n);
        const uint16_t *   p = (const uint16_t *) raw;
        for (size_t i = 0; i < n; i++) {
            f32[i] = ggml_bf16_to_fp32(*(const ggml_bf16_t *) &p[i]);
        }
        ggml_fp32_to_fp16_row(f32.data(), data, (int) n);
    }

    wctx->pending.push_back({ tensor, (const void *) data, n * sizeof(ggml_fp16_t), 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;
}

// 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);
}

// Read an array of uint32 KV values (returns empty vector if not found)
static std::vector<uint32_t> gf_get_array_u32(const GGUFModel & gf, const char * key) {
    int64_t idx = gguf_find_key(gf.gguf, key);
    if (idx < 0) {
        return {};
    }
    size_t                n   = gguf_get_arr_n(gf.gguf, idx);
    const uint32_t *      raw = (const uint32_t *) gguf_get_arr_data(gf.gguf, idx);
    std::vector<uint32_t> out(raw, raw + n);
    return out;
}