#include "model_loader.hpp"
#include "common.hpp"
#include "ggml.h"
#include "ggml-backend.h"
#include "ggml-alloc.h"
#include "ggml-cpu.h"
#include "gguf.h"
#include <cstring>
#include <cstdio>
#include <vector>
#include <utility>
#include <stdexcept>
#include "moz-overrides.h"
namespace pk {

int PromptCfg::resolve_index_or_throw(const std::string& target_lang) const {
    const std::string lang = target_lang.empty() ? default_lang : target_lang;
    int idx = lang_to_index(lang);
    if (idx < 0) {
        std::string sample;
        for (size_t i = 0; i < dict_keys.size() && i < 8; ++i)
            sample += (i ? ", " : "") + dict_keys[i];
        throw std::runtime_error("parakeet: unknown target_lang '" + lang +
                                 "'. Valid examples: " + sample + ", ...");
    }
    return idx;
}

static uint32_t kv_u32(gguf_context* g, const char* k, uint32_t d=0){
    int64_t id = gguf_find_key(g,k); return id<0 ? d : (uint32_t)gguf_get_val_u32(g,id);
}
static int32_t kv_i32(gguf_context* g, const char* k, int32_t d=0){
    int64_t id = gguf_find_key(g,k); return id<0 ? d : gguf_get_val_i32(g,id);
}
static std::vector<int32_t> kv_i32_arr(gguf_context* g, const char* k){
    std::vector<int32_t> out;
    int64_t id = gguf_find_key(g,k);
    if(id>=0 && gguf_get_arr_type(g,id)==GGUF_TYPE_INT32){
        size_t n = gguf_get_arr_n(g,id);
        const int32_t* a = (const int32_t*)gguf_get_arr_data(g,id);
        out.assign(a, a+n);
    }
    return out;
}
static std::vector<std::string> kv_str_arr(gguf_context* g, const char* k){
    std::vector<std::string> out;
    int64_t id = gguf_find_key(g,k);
    if(id>=0 && gguf_get_arr_type(g,id)==GGUF_TYPE_STRING){
        size_t n = gguf_get_arr_n(g,id);
        out.resize(n);
        for(size_t i=0;i<n;++i) out[i] = gguf_get_arr_str(g,id,i);
    }
    return out;
}
static float kv_f32(gguf_context* g, const char* k, float d=0){
    int64_t id = gguf_find_key(g,k); return id<0 ? d : gguf_get_val_f32(g,id);
}
static bool kv_bool(gguf_context* g, const char* k, bool d=false){
    int64_t id = gguf_find_key(g,k); return id<0 ? d : gguf_get_val_bool(g,id);
}
static std::string kv_str(gguf_context* g, const char* k, const char* d=""){
    int64_t id = gguf_find_key(g,k); return id<0 ? std::string(d) : std::string(gguf_get_val_str(g,id));
}
ModelLoader::~ModelLoader(){
    // Free the weight buffer BEFORE the ctxs. For the CPU path it is a from_ptr
    // buffer (free_buffer == NULL) that does NOT own its memory (ctx_ does); for
    // the device path it OWNS the device buffer. ggml_backend_buffer_free handles
    // both. device_ctx_ holds the device tensor metadata (no_alloc); ctx_ holds
    // the host source data.
    if(weights_buf_) ggml_backend_buffer_free(weights_buf_);
    if(device_ctx_) ggml_free(device_ctx_);
    if(gguf_) gguf_free(gguf_); if(ctx_) ggml_free(ctx_);
}
bool ModelLoader::realize_weights(ggml_backend_t backend){
    if(weights_buf_) return true;                       // idempotent
    if(!backend || !ctx_){ PK_LOG("realize_weights: null backend/ctx"); return false; }

    if (ggml_backend_is_cpu(backend)) {
        // Fast path: borrow the host ctx memory directly (no copy).
        // The GGUF is loaded with no_alloc=false, so every tensor's data lives
        // in one contiguous ctx mem_buffer. Wrap that exact memory as a CPU
        // backend buffer (zero-copy: ggml_backend_cpu_buffer_from_ptr borrows
        // the ptr) and point every tensor's ->buffer at it, so graphs can
        // reference the loader tensors DIRECTLY as leaves (the gallocr treats
        // data!=NULL tensors as already-allocated and never copies them;
        // reshapes/views resolve at build time). Eliminates per-call recopy.
        void*  base = ggml_get_mem_buffer(ctx_);
        size_t size = ggml_get_mem_size(ctx_);
        weights_buf_ = ggml_backend_cpu_buffer_from_ptr(base, size);
        if(!weights_buf_){ PK_LOG("realize_weights: buffer_from_ptr failed"); return false; }
        for(auto& kv : tensors_) kv.second->buffer = weights_buf_;
        return true;
    }

    // Device path (CUDA/Metal/Vulkan/...): weights must live in a backend buffer.
    // ctx_ was created no_alloc=false (host-resident data), which
    // ggml_backend_alloc_ctx_tensors rejects (it asserts the ctx is no_alloc).
    // So mirror every weight into a no_alloc=true ctx, allocate THAT on the
    // backend, upload each tensor's bytes from the host source, and repoint the
    // name->tensor map at the device tensors. ctx_ stays alive as the host source.
    const size_t n = tensors_.size();
    struct ggml_init_params dp = {
        /*.mem_size  =*/ ggml_tensor_overhead() * (n + 8),
        /*.mem_buffer=*/ nullptr,
        /*.no_alloc  =*/ true,
    };
    device_ctx_ = ggml_init(dp);
    if(!device_ctx_){ PK_LOG("realize_weights: device ctx init failed"); return false; }

    std::vector<std::pair<ggml_tensor*, const void*>> ups; ups.reserve(n);
    std::unordered_map<std::string, ggml_tensor*> devmap; devmap.reserve(n);
    for (auto& kv : tensors_) {
        ggml_tensor* s = kv.second;
        ggml_tensor* d = ggml_new_tensor(device_ctx_, s->type, GGML_MAX_DIMS, s->ne);
        ggml_set_name(d, kv.first.c_str());
        devmap.emplace(kv.first, d);
        ups.emplace_back(d, s->data);   // host source (valid in ctx_ mem buffer)
    }
    weights_buf_ = ggml_backend_alloc_ctx_tensors(device_ctx_, backend);
    if(!weights_buf_){ PK_LOG("realize_weights: alloc_ctx_tensors failed"); return false; }
    for (auto& pr : ups)
        ggml_backend_tensor_set(pr.first, pr.second, 0, ggml_nbytes(pr.first));
    tensors_.swap(devmap);   // graphs now reference the device-resident tensors
    return true;
}
bool ModelLoader::load(const std::string& path){
    struct gguf_init_params p{ /*no_alloc*/false, /*ctx*/&ctx_ };
    gguf_ = gguf_init_from_file(path.c_str(), p);
    if(!gguf_){ PK_LOG("gguf open failed: %s", path.c_str()); return false; }
    return parse_gguf();
}

// Firefox-local addition (upstreamable): load the GGUF from an already-open
// file descriptor. The sandboxed HWInference utility process is handed an fd by
// the broker and cannot open paths itself. The fd stays owned by the caller.
bool ModelLoader::load_fd(int fd){
    struct gguf_init_params p{ /*no_alloc*/false, /*ctx*/&ctx_ };
    FILE* f = fdopen(fd, "rb");
    if(!f){ PK_LOG("fdopen failed for parakeet model fd"); return false; }
    // gguf reads the whole file (no_alloc=false). Do NOT fclose(f): that would
    // close the caller-owned fd.
    gguf_ = gguf_init_from_file_ptr(f, p);
    if(!gguf_){ PK_LOG("gguf open from fd failed"); return false; }
    return parse_gguf();
}

bool ModelLoader::parse_gguf(){
    cfg_.arch        = kv_str(gguf_, "parakeet.arch");
    cfg_.feat_in     = kv_u32(gguf_, "parakeet.encoder.feat_in");
    cfg_.d_model     = kv_u32(gguf_, "parakeet.encoder.d_model");
    cfg_.n_layers    = kv_u32(gguf_, "parakeet.encoder.n_layers");
    cfg_.n_heads     = kv_u32(gguf_, "parakeet.encoder.n_heads");
    cfg_.ff_dim      = kv_u32(gguf_, "parakeet.encoder.ff_dim");
    cfg_.conv_kernel = kv_u32(gguf_, "parakeet.encoder.conv_kernel");
    cfg_.conv_norm_type = kv_str(gguf_, "parakeet.encoder.conv_norm_type", "batch_norm");
    cfg_.subsampling_factor = kv_u32(gguf_, "parakeet.encoder.subsampling_factor");
    cfg_.subsampling_conv_channels = kv_u32(gguf_, "parakeet.encoder.subsampling_conv_channels");
    cfg_.xscaling    = kv_bool(gguf_, "parakeet.encoder.xscaling", true);
    cfg_.pos_emb_max_len = kv_u32(gguf_, "parakeet.encoder.pos_emb_max_len", 5000);
    // cache-aware streaming / causal config (Phase 5). Absent for offline models
    // -> offline-safe defaults (regular style, no causal, streaming.present=false).
    cfg_.att_context_left  = kv_i32(gguf_, "parakeet.encoder.att_context_left", -1);
    cfg_.att_context_right = kv_i32(gguf_, "parakeet.encoder.att_context_right", -1);
    cfg_.att_context_style = kv_str(gguf_, "parakeet.encoder.att_context_style", "regular");
    cfg_.causal_downsampling = kv_bool(gguf_, "parakeet.encoder.causal_downsampling", false);
    cfg_.conv_causal = kv_bool(gguf_, "parakeet.encoder.conv_causal", false);
    // encoder.use_bias: false for nemotron (the attention/FFN linear projections
    // carry no bias tensor). Defaults true so existing models are unaffected.
    cfg_.use_bias = kv_bool(gguf_, "parakeet.encoder.use_bias", true);
    // Prompt conditioning (multilingual nemotron). Orthogonal capability flag;
    // absent -> present=false and the engine skips the prompt stage entirely.
    cfg_.prompt.present = kv_bool(gguf_, "parakeet.prompt.present", false);
    if(cfg_.prompt.present){
        cfg_.prompt.num_prompts  = kv_u32(gguf_, "parakeet.prompt.num_prompts", 0);
        cfg_.prompt.default_lang = kv_str(gguf_, "parakeet.prompt.default_lang", "");
        cfg_.prompt.dict_keys = kv_str_arr(gguf_, "parakeet.prompt.dictionary.keys");
        cfg_.prompt.dict_vals = kv_i32_arr(gguf_, "parakeet.prompt.dictionary.values");
    }
    if(cfg_.att_context_style != "regular"){
        StreamingCfg& s = cfg_.streaming;
        s.chunk_size = kv_i32_arr(gguf_, "parakeet.streaming.chunk_size");
        s.shift_size = kv_i32_arr(gguf_, "parakeet.streaming.shift_size");
        s.pre_encode_cache_size = kv_i32_arr(gguf_, "parakeet.streaming.pre_encode_cache_size");
        s.cache_drop_size = kv_i32(gguf_, "parakeet.streaming.cache_drop_size", 0);
        s.last_channel_cache_size = kv_i32(gguf_, "parakeet.streaming.last_channel_cache_size", 0);
        s.valid_out_len = kv_i32(gguf_, "parakeet.streaming.valid_out_len", 0);
        s.drop_extra_pre_encoded = kv_i32(gguf_, "parakeet.streaming.drop_extra_pre_encoded", 0);
        s.present = true;
    }
    cfg_.sample_rate = kv_u32(gguf_, "parakeet.preprocessor.sample_rate", 16000);
    cfg_.n_mels      = kv_u32(gguf_, "parakeet.preprocessor.n_mels");
    cfg_.n_fft       = kv_u32(gguf_, "parakeet.preprocessor.n_fft");
    cfg_.win_length  = kv_u32(gguf_, "parakeet.preprocessor.win_length");
    cfg_.hop_length  = kv_u32(gguf_, "parakeet.preprocessor.hop_length");
    cfg_.preemph     = kv_f32(gguf_, "parakeet.preprocessor.preemph", 0.0f);
    cfg_.mag_power   = kv_f32(gguf_, "parakeet.preprocessor.mag_power", 2.0f);
    cfg_.normalize   = kv_str(gguf_, "parakeet.preprocessor.normalize", "per_feature");
    cfg_.log_zero_guard = kv_f32(gguf_, "parakeet.preprocessor.log_zero_guard", 0.0f);
    cfg_.pred_hidden = kv_u32(gguf_, "parakeet.decoder.pred_hidden");
    cfg_.pred_rnn_layers = kv_u32(gguf_, "parakeet.decoder.pred_rnn_layers");
    cfg_.joint_hidden = kv_u32(gguf_, "parakeet.joint.joint_hidden");
    cfg_.joint_activation = kv_str(gguf_, "parakeet.joint.activation");
    cfg_.max_symbols = kv_u32(gguf_, "parakeet.decoding.max_symbols", 10);
    cfg_.vocab_size  = kv_u32(gguf_, "parakeet.vocab_size");
    cfg_.blank_id    = kv_u32(gguf_, "parakeet.blank_id");
    // durations array (stored as INT32 by the converter)
    { int64_t id = gguf_find_key(gguf_, "parakeet.tdt.durations");
      if(id>=0 && gguf_get_arr_type(gguf_,id)==GGUF_TYPE_INT32){
          size_t n = gguf_get_arr_n(gguf_,id);
          const int32_t* a = (const int32_t*)gguf_get_arr_data(gguf_,id);
          cfg_.tdt_durations.assign(a, a+n); } }
    // tokenizer pieces STRING array
    { int64_t id = gguf_find_key(gguf_, "parakeet.tokenizer.pieces");
      if(id>=0 && gguf_get_arr_type(gguf_,id)==GGUF_TYPE_STRING){
          size_t n = gguf_get_arr_n(gguf_,id);
          cfg_.tokenizer_pieces.resize(n);
          for(size_t i=0;i<n;++i)
              cfg_.tokenizer_pieces[i] = gguf_get_arr_str(gguf_,id,i); } }
    // tensors
    const int64_t nt = gguf_get_n_tensors(gguf_);
    for(int64_t i=0;i<nt;++i){ const char* nm = gguf_get_tensor_name(gguf_,i);
        ggml_tensor* t = ggml_get_tensor(ctx_, nm); if(t) tensors_[nm]=t; }
    return cfg_.d_model>0 && cfg_.vocab_size>0;
}
ggml_tensor* ModelLoader::tensor(const std::string& n) const {
    auto it = tensors_.find(n); return it==tensors_.end()? nullptr : it->second;
}
}