/**
 * Meg v1 - Native Engine Coordinator Implementation with Phonetic N-Gram Stitching
 */

#include "meg_v1_engine.hpp"
#include <fstream>
#include <sstream>
#include <iostream>
#include <algorithm>
#include <cmath>
#include <cstring>

namespace megv1 {

static std::vector<std::string> split_words(const std::string& str) {
    std::vector<std::string> words;
    std::istringstream iss(str);
    std::string w;
    while (iss >> w) {
        words.push_back(w);
    }
    return words;
}

static std::string trim(const std::string& str) {
    size_t first = str.find_first_not_of(" \t\n\r");
    if (first == std::string::npos) return "";
    size_t last = str.find_last_not_of(" \t\n\r");
    return str.substr(first, (last - first + 1));
}

static std::string to_lower_str(const std::string& s) {
    std::string r = s;
    std::transform(r.begin(), r.end(), r.begin(), [](unsigned char c) { return std::tolower(c); });
    return r;
}

static const std::unordered_set<std::string> TIME_MARKERS = {"AM", "PM", "am", "pm", "A.M.", "P.M."};
static const std::unordered_set<std::string> PREPOSITIONS_POST_TIME = {
    "with", "for", "on", "in", "at", "near", "to", "by", "before", "after", "and"
};
static const std::unordered_set<std::string> NUMBER_WORDS = {
    "one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten"
};
static const std::unordered_set<std::string> TRANSITIVE_VERBS = {
    "export", "set", "check", "clone", "run", "build", "dispatch", "test",
    "verify", "create", "configure", "start", "stop", "restart", "deploy", "call", "load", "import"
};

MegV1Engine::MegV1Engine(const megv1_config_t& config) : config_(config) {
    init_guardrails();
    if (config.vocab_path && strlen(config.vocab_path) > 0) {
        init_tagger(config.tagger_model_path, config.vocab_path);
    }
    if (config.asr_model_dir && strlen(config.asr_model_dir) > 0) {
        init_asr(config.asr_model_dir);
    }
}

MegV1Engine::~MegV1Engine() {
    if (asr_stream_ && asr_recognizer_) {
        SherpaOnnxDestroyOnlineStream(asr_stream_);
        asr_stream_ = nullptr;
    }
    if (asr_recognizer_) {
        SherpaOnnxDestroyOnlineRecognizer(asr_recognizer_);
        asr_recognizer_ = nullptr;
    }
}

void MegV1Engine::init_guardrails() {
    // Canonical tech terms
    const std::vector<std::pair<std::string, std::string>> tech_list = {
        {"macos", "macOS"}, {"ios", "iOS"}, {"ipados", "iPadOS"}, {"watchos", "watchOS"},
        {"swiftui", "SwiftUI"}, {"appkit", "AppKit"}, {"uikit", "UIKit"},
        {"pytorch", "PyTorch"}, {"onnx", "ONNX"}, {"graphql", "GraphQL"},
        {"kubernetes", "Kubernetes"}, {"docker", "Docker"}, {"github", "GitHub"},
        {"cgevent", "CGEvent"}, {"axisprocesstrusted", "AXIsProcessTrusted"},
        {"api", "API"}, {"sdk", "SDK"}, {"cpu", "CPU"}, {"gpu", "GPU"},
        {"mps", "MPS"}, {"int8", "INT8"}, {"fp16", "FP16"}, {"fp32", "FP32"},
        {"dear machine", "Dear Machine"}
    };

    for (const auto& kv : tech_list) {
        canonical_tech_terms_[to_lower_str(kv.first)] = kv.second;
        hotword_trie_.insert(kv.first);
    }

    // Regex invariants
    protected_regexes_.push_back(std::regex(R"(^[a-z]+(?:[A-Z][a-z0-9]*)+$)")); // camelCase
    protected_regexes_.push_back(std::regex(R"(^[a-z0-9]+(?:_[a-z0-9]+)+$)")); // snake_case
    protected_regexes_.push_back(std::regex(R"(^\$?\d+(?:,\d{3})*(?:\.\d+)?%?$)")); // Numbers / Currency
    protected_regexes_.push_back(std::regex(R"(^https?:\/\/[^\s]+$)")); // URLs
    protected_regexes_.push_back(std::regex(R"(^\d{1,2}:\d{2}(?::\d{2})?(?:[aApP][mM])?$)")); // Time
}

void MegV1Engine::init_asr(const std::string& asr_model_dir) {
    SherpaOnnxOnlineRecognizerConfig asr_config;
    std::memset(&asr_config, 0, sizeof(asr_config));

    std::string encoder = asr_model_dir + "/encoder-epoch-99-avg-1.int8.onnx";
    std::string decoder = asr_model_dir + "/decoder-epoch-99-avg-1.int8.onnx";
    std::string joiner = asr_model_dir + "/joiner-epoch-99-avg-1.int8.onnx";
    std::string tokens = asr_model_dir + "/tokens.txt";

    asr_config.feat_config.sample_rate = config_.sample_rate > 0 ? config_.sample_rate : 16000;
    asr_config.feat_config.feature_dim = 80;

    asr_config.model_config.transducer.encoder = encoder.c_str();
    asr_config.model_config.transducer.decoder = decoder.c_str();
    asr_config.model_config.transducer.joiner = joiner.c_str();
    asr_config.model_config.tokens = tokens.c_str();
    asr_config.model_config.num_threads = config_.num_threads > 0 ? config_.num_threads : 2;
    asr_config.model_config.debug = 0;
    asr_config.model_config.provider = "cpu";

    asr_config.decoding_method = "greedy_search";
    asr_config.max_active_paths = 4;
    asr_config.enable_endpoint = 1;
    asr_config.rule1_min_trailing_silence = 2.4f;
    asr_config.rule2_min_trailing_silence = 0.8f;
    asr_config.rule3_min_utterance_length = 20.0f;

    // Contextual Biasing Hotwords
    static std::string hotwords_str = (
        "AXIS PROCESS TRUSTED/15.0\n"
        "CG EVENT/12.0\n"
        "PYTORCH/12.0\n"
        "PI TORCH/12.0\n"
        "ONNX/15.0\n"
        "MACOS/10.0\n"
        "MAC OS/10.0\n"
        "DEAR MACHINE/15.0\n"
        "SWIFTUI/12.0\n"
        "GRAPHQL/12.0\n"
        "KUBERNETES/10.0\n"
        "DOCKER/10.0\n"
    );
    asr_config.hotwords_buf = hotwords_str.c_str();
    asr_config.hotwords_buf_size = static_cast<int32_t>(hotwords_str.size());
    asr_config.hotwords_score = 5.0f;

    asr_recognizer_ = SherpaOnnxCreateOnlineRecognizer(&asr_config);
    if (asr_recognizer_) {
        asr_stream_ = SherpaOnnxCreateOnlineStream(asr_recognizer_);
    }
}

void MegV1Engine::init_tagger(const std::string& tagger_path, const std::string& vocab_path) {
    // Load WordPiece Vocab
    std::ifstream vf(vocab_path);
    if (vf.is_open()) {
        std::string line;
        int64_t id = 0;
        while (std::getline(vf, line)) {
            line = trim(line);
            if (!line.empty()) {
                vocab_to_id_[line] = id;
                id_to_vocab_[id] = line;
                id++;
            }
        }
        if (vocab_to_id_.find("[UNK]") != vocab_to_id_.end()) unk_id_ = vocab_to_id_["[UNK]"];
        if (vocab_to_id_.find("[CLS]") != vocab_to_id_.end()) cls_id_ = vocab_to_id_["[CLS]"];
        if (vocab_to_id_.find("[SEP]") != vocab_to_id_.end()) sep_id_ = vocab_to_id_["[SEP]"];
        if (vocab_to_id_.find("[PAD]") != vocab_to_id_.end()) pad_id_ = vocab_to_id_["[PAD]"];
    }

    // Init ONNX Runtime
    ort_env_ = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "MegV1Tagger");
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(config_.num_threads > 0 ? config_.num_threads : 2);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

    ort_session_ = std::make_unique<Ort::Session>(*ort_env_, tagger_path.c_str(), session_options);
    ort_memory_info_ = std::make_unique<Ort::MemoryInfo>(Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault));

    input_node_names_ = {"input_ids", "attention_mask"};
    output_node_names_ = {"action_logits", "punct_logits", "casing_logits"};
}

std::vector<int64_t> MegV1Engine::tokenize_words(const std::vector<std::string>& words, std::vector<int>& out_first_subword_indices) {
    std::vector<int64_t> input_ids;
    out_first_subword_indices.clear();
    input_ids.push_back(cls_id_);

    for (size_t w_idx = 0; w_idx < words.size(); ++w_idx) {
        std::string w = to_lower_str(words[w_idx]);
        out_first_subword_indices.push_back(static_cast<int>(input_ids.size()));

        if (vocab_to_id_.find(w) != vocab_to_id_.end()) {
            input_ids.push_back(vocab_to_id_[w]);
        } else {
            bool is_bad = false;
            size_t start = 0;
            std::vector<int64_t> sub_tokens;
            while (start < w.length()) {
                size_t end = w.length();
                std::string cur_substr;
                int64_t cur_id = -1;
                while (start < end) {
                    std::string substr = w.substr(start, end - start);
                    if (start > 0) substr = "##" + substr;
                    if (vocab_to_id_.find(substr) != vocab_to_id_.end()) {
                        cur_substr = substr;
                        cur_id = vocab_to_id_[substr];
                        break;
                    }
                    end--;
                }
                if (cur_id == -1) {
                    is_bad = true;
                    break;
                }
                sub_tokens.push_back(cur_id);
                start = end;
            }
            if (is_bad || sub_tokens.empty()) {
                input_ids.push_back(unk_id_);
            } else {
                input_ids.insert(input_ids.end(), sub_tokens.begin(), sub_tokens.end());
            }
        }
    }
    input_ids.push_back(sep_id_);
    return input_ids;
}

bool MegV1Engine::is_protected(const std::string& word) const {
    std::string clean = word;
    while (!clean.empty() && (clean.back() == '.' || clean.back() == ',' || clean.back() == '?' || clean.back() == '!' || clean.back() == ':')) {
        clean.pop_back();
    }
    if (clean.empty()) return false;

    // Check Hotword Trie
    if (hotword_trie_.contains(clean)) return true;

    // Check Canonical terms
    std::string lower = to_lower_str(clean);
    if (canonical_tech_terms_.find(lower) != canonical_tech_terms_.end()) return true;

    // Check Regex invariants
    for (const auto& rx : protected_regexes_) {
        if (std::regex_match(clean, rx)) return true;
    }
    return false;
}

bool MegV1Engine::is_stutter_or_repaired(const std::vector<std::string>& words, size_t idx) const {
    std::string word_clean = to_lower_str(words[idx]);
    size_t max_ahead = std::min(words.size(), idx + 6);
    
    // Stutter duplicate
    for (size_t i = idx + 1; i < max_ahead; ++i) {
        if (to_lower_str(words[i]) == word_clean) return true;
    }

    // Speech repair cues
    static const std::vector<std::string> cues = {
        "or rather", "wait no", "scratch that", "actually no", "sorry", "i mean", "make that"
    };

    std::string window_str = "";
    for (size_t i = idx + 1; i < max_ahead; ++i) {
        if (!window_str.empty()) window_str += " ";
        window_str += to_lower_str(words[i]);
    }

    for (const auto& cue : cues) {
        if (window_str.find(cue) == 0 || window_str.find(" " + cue) != std::string::npos) {
            return true;
        }
    }
    return false;
}

std::string MegV1Engine::render_slice(const std::vector<std::string>& words, bool is_final_boundary) {
    if (words.empty()) return "";
    if (!ort_session_) {
        std::string res;
        for (const auto& w : words) {
            if (!res.empty()) res += " ";
            res += w;
        }
        return res;
    }

    std::vector<int> first_subword_indices;
    std::vector<int64_t> input_ids = tokenize_words(words, first_subword_indices);
    size_t seq_len = input_ids.size();
    std::vector<int64_t> attention_mask(seq_len, 1);

    std::vector<int64_t> input_shape = {1, static_cast<int64_t>(seq_len)};
    std::vector<Ort::Value> input_tensors;
    input_tensors.push_back(Ort::Value::CreateTensor<int64_t>(
        *ort_memory_info_, input_ids.data(), seq_len, input_shape.data(), input_shape.size()
    ));
    input_tensors.push_back(Ort::Value::CreateTensor<int64_t>(
        *ort_memory_info_, attention_mask.data(), seq_len, input_shape.data(), input_shape.size()
    ));

    const char* in_names[] = {"input_ids", "attention_mask"};
    const char* out_names[] = {"action_logits", "punct_logits", "casing_logits"};

    auto output_tensors = ort_session_->Run(
        Ort::RunOptions{nullptr}, in_names, input_tensors.data(), input_tensors.size(), out_names, 3
    );

    const float* act_logits = output_tensors[0].GetTensorData<float>();
    const float* punc_logits = output_tensors[1].GetTensorData<float>();
    const float* case_logits = output_tensors[2].GetTensorData<float>();

    // Decode predictions per word
    std::vector<std::string> rendered_words;
    bool capitalize_next = true;

    for (size_t w_idx = 0; w_idx < words.size(); ++w_idx) {
        int token_pos = first_subword_indices[w_idx];
        const std::string& orig_word = words[w_idx];

        int act_pred = (act_logits[token_pos * 2 + 1] > act_logits[token_pos * 2]) ? 1 : 0; // 0: KEEP, 1: DELETE

        int punc_pred = 0;
        float max_p = punc_logits[token_pos * 5];
        for (int p = 1; p < 5; ++p) {
            if (punc_logits[token_pos * 5 + p] > max_p) {
                max_p = punc_logits[token_pos * 5 + p];
                punc_pred = p;
            }
        }

        int case_pred = 0;
        float max_c = case_logits[token_pos * 3];
        for (int c = 1; c < 3; ++c) {
            if (case_logits[token_pos * 3 + c] > max_c) {
                max_c = case_logits[token_pos * 3 + c];
                case_pred = c;
            }
        }

        bool is_prot = is_protected(orig_word);

        // Guardrail Invariant Check: Protected entities NEVER deleted unless legitimate repair
        if (config_.enable_guardrails && is_prot) {
            bool is_stutter_repair = is_stutter_or_repaired(words, w_idx);
            if (act_pred == 1 && !is_stutter_repair) {
                act_pred = 0; // Force KEEP
            }
        }

        if (act_pred == 1) {
            continue; // DELETE
        }

        std::string clean_base = orig_word;
        while (!clean_base.empty() && (clean_base.back() == '.' || clean_base.back() == ',' || clean_base.back() == '?' || clean_base.back() == '!' || clean_base.back() == ':')) {
            clean_base.pop_back();
        }
        std::string lower_clean = to_lower_str(clean_base);

        // Punctuation symbol
        std::string punct_sym = "";
        if (punc_pred == 1) punct_sym = ".";
        else if (punc_pred == 2) punct_sym = ",";
        else if (punc_pred == 3) punct_sym = "?";
        else if (punc_pred == 4) punct_sym = ":";

        // Post-timestamp preposition punctuation fix
        if (TIME_MARKERS.find(clean_base) != TIME_MARKERS.end()) {
            if (w_idx + 1 < words.size()) {
                std::string next_w = to_lower_str(words[w_idx + 1]);
                if (PREPOSITIONS_POST_TIME.find(next_w) != PREPOSITIONS_POST_TIME.end()) {
                    punct_sym = "";
                }
            }
        }

        // Transitive verb lookahead guard: suppress period if followed by direct object or not final boundary
        if (TRANSITIVE_VERBS.find(lower_clean) != TRANSITIVE_VERBS.end()) {
            if (!is_final_boundary || (w_idx + 1 < words.size())) {
                if (punct_sym == ".") punct_sym = "";
            }
        }

        // Casing & Entity Invariants
        bool has_lower = false;
        bool has_up = false;
        for (char c : clean_base) {
            if (std::islower(static_cast<unsigned char>(c))) has_lower = true;
            if (std::isupper(static_cast<unsigned char>(c))) has_up = true;
        }
        bool is_mixed_case = (has_lower && has_up);

        std::string rendered;
        if (is_prot || is_mixed_case) {
            std::string canon;
            if (canonical_tech_terms_.find(lower_clean) != canonical_tech_terms_.end()) {
                rendered = canonical_tech_terms_[lower_clean];
            } else if (hotword_trie_.contains(clean_base, &canon)) {
                rendered = canon;
            } else {
                rendered = clean_base;
            }
        } else {
            if (case_pred == 2) { // UPPER
                rendered = clean_base;
                std::transform(rendered.begin(), rendered.end(), rendered.begin(), [](unsigned char c) { return std::toupper(c); });
            } else if (capitalize_next || case_pred == 1) { // TITLE
                rendered = lower_clean;
                if (!rendered.empty()) rendered[0] = std::toupper(rendered[0]);
            } else if (NUMBER_WORDS.find(lower_clean) != NUMBER_WORDS.end()) {
                rendered = lower_clean;
            } else {
                rendered = lower_clean;
            }
        }

        if (capitalize_next && !is_prot && !rendered.empty()) {
            rendered[0] = std::toupper(rendered[0]);
            capitalize_next = false;
        }

        if (punct_sym == "." || punct_sym == "?") {
            capitalize_next = true;
        }

        rendered_words.push_back(rendered + punct_sym);
    }

    if (rendered_words.empty()) return "";

    if (is_final_boundary) {
        std::string& last = rendered_words.back();
        if (last.back() == ',' || last.back() == ':') {
            last.back() = '.';
        } else if (last.back() != '.' && last.back() != '?' && last.back() != '!') {
            last += ".";
        }
    }

    std::string out_str;
    for (size_t i = 0; i < rendered_words.size(); ++i) {
        if (i > 0) out_str += " ";
        out_str += rendered_words[i];
    }
    return out_str;
}

void MegV1Engine::process_hypothesis(const std::string& hypothesis, bool is_endpoint) {
    std::vector<std::string> raw_tokens = split_words(hypothesis);
    if (raw_tokens.empty()) return;

    // Step 1: Fragmented Disfluency & Suffix Filter
    // Filter leading orphaned suffix fragments (e.g. "ly" from partial "basically")
    std::vector<std::string> filtered_tokens;
    for (size_t i = 0; i < raw_tokens.size(); ++i) {
        std::string norm = to_lower_str(raw_tokens[i]);
        if (i == 0 && (norm == "ly" || norm == "er" || norm == "ah" || norm == "um" || norm == "uh")) {
            continue; // Drop leading orphaned acoustic disfluency fragment
        }
        filtered_tokens.push_back(raw_tokens[i]);
    }

    // Step 2: Multi-Token Phonetic N-Gram Stitcher
    std::vector<std::string> tokens = hotword_trie_.stitch_ngrams(filtered_tokens);
    size_t total_tokens = tokens.size();
    size_t lookahead = config_.lookahead_window > 0 ? config_.lookahead_window : 3;

    size_t frontier_idx = is_endpoint ? total_tokens : (total_tokens > lookahead ? total_tokens - lookahead : 0);

    // Step 3: Commit stable tokens
    if (frontier_idx > committed_count_in_hypothesis_) {
        std::vector<std::string> new_commit_slice(
            tokens.begin() + committed_count_in_hypothesis_, tokens.begin() + frontier_idx
        );

        std::vector<std::string> slice_with_prefix;
        size_t prefix_len = std::min(committed_words_.size(), static_cast<size_t>(2));
        if (prefix_len > 0) {
            slice_with_prefix.insert(
                slice_with_prefix.end(), committed_words_.end() - prefix_len, committed_words_.end()
            );
        }
        slice_with_prefix.insert(slice_with_prefix.end(), new_commit_slice.begin(), new_commit_slice.end());

        std::string full_render = render_slice(slice_with_prefix, is_endpoint);
        std::string commit_chunk;

        if (prefix_len > 0) {
            std::vector<std::string> prefix_words(committed_words_.end() - prefix_len, committed_words_.end());
            std::string prefix_render = render_slice(prefix_words, false);
            if (full_render.length() >= prefix_render.length() && full_render.substr(0, prefix_render.length()) == prefix_render) {
                commit_chunk = trim(full_render.substr(prefix_render.length()));
            } else {
                std::vector<std::string> rendered_all = split_words(full_render);
                std::string fallback;
                for (size_t i = prefix_len; i < rendered_all.size(); ++i) {
                    if (!fallback.empty()) fallback += " ";
                    fallback += rendered_all[i];
                }
                commit_chunk = fallback;
            }
        } else {
            commit_chunk = full_render;
        }

        if (!commit_chunk.empty()) {
            committed_clean_chunks_.push_back(commit_chunk);
            committed_words_.insert(committed_words_.end(), new_commit_slice.begin(), new_commit_slice.end());

            if (committed_cb_) {
                committed_cb_(commit_chunk.c_str(), user_data_);
            }
        }
        committed_count_in_hypothesis_ = frontier_idx;
    }

    // Step 4: Active Tail Preview
    std::vector<std::string> active_tail(tokens.begin() + frontier_idx, tokens.end());
    if (!active_tail.empty()) {
        std::vector<std::string> slice_with_prefix;
        size_t prefix_len = std::min(committed_words_.size(), static_cast<size_t>(2));
        if (prefix_len > 0) {
            slice_with_prefix.insert(
                slice_with_prefix.end(), committed_words_.end() - prefix_len, committed_words_.end()
            );
        }
        slice_with_prefix.insert(slice_with_prefix.end(), active_tail.begin(), active_tail.end());

        std::string full_render = render_slice(slice_with_prefix, is_endpoint);
        std::string partial_clean;

        if (prefix_len > 0) {
            std::vector<std::string> prefix_words(committed_words_.end() - prefix_len, committed_words_.end());
            std::string prefix_render = render_slice(prefix_words, false);
            if (full_render.length() >= prefix_render.length() && full_render.substr(0, prefix_render.length()) == prefix_render) {
                partial_clean = trim(full_render.substr(prefix_render.length()));
            } else {
                std::vector<std::string> rendered_all = split_words(full_render);
                std::string fallback;
                for (size_t i = prefix_len; i < rendered_all.size(); ++i) {
                    if (!fallback.empty()) fallback += " ";
                    fallback += rendered_all[i];
                }
                partial_clean = fallback;
            }
        } else {
            partial_clean = full_render;
        }

        last_partial_clean_ = partial_clean;
    } else {
        last_partial_clean_ = "";
    }

    if (partial_cb_) {
        std::string tail_str;
        for (const auto& w : active_tail) {
            if (!tail_str.empty()) tail_str += " ";
            tail_str += w;
        }
        partial_cb_(last_partial_clean_.c_str(), tail_str.c_str(), user_data_);
    }

    if (is_endpoint) {
        committed_count_in_hypothesis_ = 0;
        last_hypothesis_ = "";
        last_partial_clean_ = "";
    }
}

void MegV1Engine::set_callbacks(megv1_partial_cb_t partial_cb, megv1_committed_cb_t committed_cb, void* user_data) {
    std::lock_guard<std::mutex> lock(engine_mutex_);
    partial_cb_ = partial_cb;
    committed_cb_ = committed_cb;
    user_data_ = user_data;
}

void MegV1Engine::feed_pcm(const float* samples, int32_t num_samples) {
    std::lock_guard<std::mutex> lock(engine_mutex_);
    if (!asr_recognizer_ || !asr_stream_ || !samples || num_samples <= 0) return;

    SherpaOnnxOnlineStreamAcceptWaveform(asr_stream_, config_.sample_rate > 0 ? config_.sample_rate : 16000, samples, num_samples);

    while (SherpaOnnxIsOnlineStreamReady(asr_recognizer_, asr_stream_)) {
        SherpaOnnxDecodeOnlineStream(asr_recognizer_, asr_stream_);
    }

    int is_endpoint = SherpaOnnxOnlineStreamIsEndpoint(asr_recognizer_, asr_stream_);
    const SherpaOnnxOnlineRecognizerResult* r = SherpaOnnxGetOnlineStreamResult(asr_recognizer_, asr_stream_);

    if (r && r->text) {
        std::string hyp = trim(r->text);
        process_hypothesis(hyp, is_endpoint != 0);
        SherpaOnnxDestroyOnlineRecognizerResult(r);
    }

    if (is_endpoint != 0) {
        SherpaOnnxOnlineStreamReset(asr_recognizer_, asr_stream_);
    }
}

void MegV1Engine::add_hotwords(const char** words, int32_t num_words) {
    std::lock_guard<std::mutex> lock(engine_mutex_);
    if (!words || num_words <= 0) return;

    for (int32_t i = 0; i < num_words; ++i) {
        if (words[i]) {
            hotword_trie_.insert(words[i]);
        }
    }
}

void MegV1Engine::reset() {
    std::lock_guard<std::mutex> lock(engine_mutex_);
    if (asr_recognizer_ && asr_stream_) {
        SherpaOnnxOnlineStreamReset(asr_recognizer_, asr_stream_);
    }
    committed_words_.clear();
    committed_clean_chunks_.clear();
    last_hypothesis_.clear();
    committed_count_in_hypothesis_ = 0;
    last_partial_clean_.clear();
    cached_full_text_.clear();
}

std::string MegV1Engine::get_full_text() const {
    std::lock_guard<std::mutex> lock(engine_mutex_);
    std::string res;
    for (const auto& chunk : committed_clean_chunks_) {
        if (!chunk.empty()) {
            if (!res.empty()) res += " ";
            res += chunk;
        }
    }
    return res;
}

} // namespace megv1

// C ABI Implementation

extern "C" {

struct megv1_handle {
    std::unique_ptr<megv1::MegV1Engine> impl;
    std::string cached_str;
};

struct megv1_engine {
    std::unique_ptr<megv1::MegV1Engine> impl;
    std::string cached_str;
};
typedef struct megv1_engine megv1_engine_t;

megv1_t* megv1_create(const char* model_dir) {
    try {
        std::string base_dir = (model_dir && strlen(model_dir) > 0) ? model_dir : "models";
        
        megv1_config_t cfg;
        std::string asr_dir = base_dir + "/asr";
        std::string tagger_path = base_dir + "/tagger/meg_v1_tagger_int8.onnx";
        std::string vocab_path = base_dir + "/tagger/vocab.txt";

        // Check fallback flat structure or models/ structure
        std::ifstream test_f(tagger_path);
        if (!test_f.good()) {
            asr_dir = base_dir;
            tagger_path = base_dir + "/meg_v1_tagger_int8.onnx";
            vocab_path = base_dir + "/vocab.txt";
        }

        cfg.asr_model_dir = asr_dir.c_str();
        cfg.tagger_model_path = tagger_path.c_str();
        cfg.vocab_path = vocab_path.c_str();
        cfg.lookahead_window = 3;
        cfg.sample_rate = 16000;
        cfg.num_threads = 2;
        cfg.enable_guardrails = true;

        auto handle = new megv1_handle();
        handle->impl = std::make_unique<megv1::MegV1Engine>(cfg);
        return handle;
    } catch (const std::exception& e) {
        std::cerr << "[meg_v1 Error] Failed to create engine: " << e.what() << std::endl;
        return nullptr;
    }
}

void megv1_set_callbacks(
    megv1_t* handle,
    megv1_partial_cb_t partial_cb,
    megv1_committed_cb_t committed_cb,
    void* user_data
) {
    if (handle && handle->impl) {
        handle->impl->set_callbacks(partial_cb, committed_cb, user_data);
    }
}

void megv1_feed_pcm(
    megv1_t* handle,
    const float* samples,
    int count
) {
    if (handle && handle->impl) {
        handle->impl->feed_pcm(samples, count);
    }
}

void megv1_add_hotwords(
    megv1_t* handle,
    const char* const* words,
    int count
) {
    if (handle && handle->impl) {
        handle->impl->add_hotwords(const_cast<const char**>(words), count);
    }
}

void megv1_reset(megv1_t* handle) {
    if (handle && handle->impl) {
        handle->impl->reset();
    }
}

const char* megv1_get_full_text(megv1_t* handle) {
    if (handle && handle->impl) {
        handle->cached_str = handle->impl->get_full_text();
        return handle->cached_str.c_str();
    }
    return "";
}

void megv1_destroy(megv1_t* handle) {
    if (handle) {
        delete handle;
    }
}

// Backwards-compatible legacy exports
megv1_config_t megv1_default_config(void) {
    megv1_config_t cfg;
    cfg.asr_model_dir = "models/asr";
    cfg.tagger_model_path = "models/tagger/meg_v1_tagger_int8.onnx";
    cfg.vocab_path = "models/tagger/vocab.txt";
    cfg.lookahead_window = 3;
    cfg.sample_rate = 16000;
    cfg.num_threads = 2;
    cfg.enable_guardrails = true;
    return cfg;
}

megv1_engine_t* megv1_engine_create(const megv1_config_t* config) {
    try {
        megv1_config_t cfg = config ? *config : megv1_default_config();
        auto handle = new megv1_engine();
        handle->impl = std::make_unique<megv1::MegV1Engine>(cfg);
        return handle;
    } catch (const std::exception& e) {
        std::cerr << "[meg_v1 C API Error] Failed to create engine: " << e.what() << std::endl;
        return nullptr;
    }
}

} // extern "C"