#pragma once
#include <string>
#include <vector>
#include <unordered_map>
#include <cstdint>
struct ggml_tensor;
struct ggml_context;
struct gguf_context;
struct ggml_backend_buffer;
typedef struct ggml_backend_buffer* ggml_backend_buffer_t;
struct ggml_backend;
typedef struct ggml_backend* ggml_backend_t;
namespace pk {
// Cache-aware streaming params (Phase 5), populated only for streaming models
// (att_context_style != "regular"). Mirrors NeMo CacheAwareStreamingConfig.
// List fields are stored as int32 arrays in the GGUF (e.g. chunk_size=[9,16]).
struct StreamingCfg {
    std::vector<int32_t> chunk_size;              // per-step encoder chunk frames
    std::vector<int32_t> shift_size;              // per-step shift frames
    std::vector<int32_t> pre_encode_cache_size;   // pre-encode cache (mel frames)
    int32_t cache_drop_size=0;
    int32_t last_channel_cache_size=0;            // attention left-context cache
    int32_t valid_out_len=0;                      // valid encoder frames per step
    int32_t drop_extra_pre_encoded=0;
    bool present=false;                           // true only for streaming models
};
// Prompt-conditioning config (multilingual nemotron). present=false for all
// existing models, which then skip the prompt stage entirely.
struct PromptCfg {
    bool present = false;
    uint32_t num_prompts = 0;
    std::string default_lang;                       // e.g. "auto"
    std::vector<std::string> dict_keys;             // locale strings
    std::vector<int32_t>     dict_vals;             // parallel prompt indices
    // Resolve a locale to its prompt index; -1 if unknown.
    int lang_to_index(const std::string& lang) const {
        for (size_t i = 0; i < dict_keys.size(); ++i)
            if (dict_keys[i] == lang) return (int)dict_vals[i];
        return -1;
    }
    // Resolve target_lang to its prompt index, applying the model default for an
    // empty string. THROWS std::runtime_error on an unknown locale. Shared by the
    // offline (Model::resolve_prompt_index) and streaming (StreamingSession ctor)
    // paths so both reject typos identically (matches the C-API contract).
    int resolve_index_or_throw(const std::string& target_lang) const;
};
struct ParakeetConfig {
    std::string arch;
    // encoder
    uint32_t feat_in=0, d_model=0, n_layers=0, n_heads=0, ff_dim=0, conv_kernel=0;
    std::string conv_norm_type;
    uint32_t subsampling_factor=0, subsampling_conv_channels=0, pos_emb_max_len=5000;
    bool xscaling=true;
    // cache-aware streaming / causal config (Phase 5; offline-safe defaults)
    int32_t att_context_left=-1, att_context_right=-1; // [-1,-1] = full context
    std::string att_context_style="regular";            // or "chunked_limited"
    bool causal_downsampling=false;                     // causal subsampling pad
    bool conv_causal=false;                             // causal depthwise conv pad
    bool use_bias=true;     // false for nemotron (encoder linears have no bias)
    StreamingCfg streaming;
    PromptCfg prompt;       // prompt conditioning (present=false for non-prompt)
    // preprocessor
    uint32_t sample_rate=16000, n_mels=0, n_fft=0, win_length=0, hop_length=0;
    float preemph=0.0f, mag_power=2.0f, log_zero_guard=0.0f;
    std::string normalize;
    // transducer (optional)
    uint32_t pred_hidden=0, pred_rnn_layers=0, joint_hidden=0;
    std::string joint_activation;
    std::vector<int32_t> tdt_durations;
    uint32_t max_symbols=10;  // greedy max symbols per frame (NeMo default 10)
    // vocab
    uint32_t vocab_size=0, blank_id=0;
    std::vector<std::string> tokenizer_pieces;
};
class ModelLoader {
public:
    ModelLoader() = default;
    ~ModelLoader();
    bool load(const std::string& path);
    // Firefox-local: load from an already-open fd (sandboxed process).
    bool load_fd(int fd);
    const ParakeetConfig& config() const { return cfg_; }
    const std::vector<std::string>& tokenizer_pieces() const { return cfg_.tokenizer_pieces; }
    ggml_tensor* tensor(const std::string& name) const; // nullptr if absent
    ggml_context* ggml_ctx() const { return ctx_; }

    // Give every weight tensor a CPU backend buffer (ONCE), so graphs can
    // reference the loader's tensors DIRECTLY as leaves with zero per-call
    // copying. The GGUF is loaded with no_alloc=false, so all weight data
    // already lives in one contiguous ctx mem_buffer; this wraps that exact
    // memory via ggml_backend_cpu_buffer_from_ptr (no data movement) and points
    // every tensor's ->buffer at it. After this the weights are valid graph
    // leaves on the CPU backend and reshapes/views of them resolve their data
    // at build time. Idempotent; safe to call once at load. The backend must be
    // the same CPU backend the compute path uses. Returns false on failure.
    bool realize_weights(ggml_backend_t backend);
    bool weights_realized() const { return weights_buf_ != nullptr; }
private:
    // Parse metadata + tensors from an already-opened gguf_. Shared by load()
    // and load_fd().
    bool parse_gguf();
    ParakeetConfig cfg_;
    gguf_context* gguf_ = nullptr;
    ggml_context* ctx_ = nullptr;
    // CPU backend: wraps ctx_ mem_buffer (zero-copy). Device backend: owns the
    // device buffer holding the uploaded weights (mirrored into device_ctx_).
    ggml_backend_buffer_t weights_buf_ = nullptr;
    ggml_context* device_ctx_ = nullptr;  // no_alloc mirror ctx for device weights
    std::unordered_map<std::string, ggml_tensor*> tensors_;
};
}