#pragma once
#include <onnxruntime_cxx_api.h>
#include <tokenizers_cpp.h>   // mlc-ai/tokenizers-cpp — reads tokenizer.json

#include <string>
#include <vector>
#include <memory>
#include <fstream>
#include <sstream>
#include <algorithm>
#include <cctype>
#include <unicode/unistr.h>   // ICU — for NFKC + case fold
#include <unicode/normalizer2.h>

// ---------------------------------------------------------------------------
// EOUDetector  —  End-of-Utterance probability via Qwen2.5-0.5B ONNX model
//
// ONNX interface:
//   input  "input_ids"  INT64  [1, seq_len]   (seq_len dynamic)
//   output "eou_prob"   FLOAT  [1]
//
// Usage:
//   EOUDetector eou("path/to/model.onnx", "path/to/tokenizer.json");
//   float p = eou.score(history, current_text);
//   if (p >= EOUDetector::THRESHOLD) { /* fire endpoint */ }
// ---------------------------------------------------------------------------

class EOUDetector {
public:
    // Tuned on 200 labelled samples: FA=22.3%, FR=6.0%, F1=0.851
    static constexpr float THRESHOLD = 0.0766f;
    static constexpr int   MAX_TOKENS = 512;

    struct Turn {
        enum class Role { User, Assistant } role;
        std::string text;   // raw text — preprocessing applied internally
    };

    // modelPath   — directory containing model.onnx AND model.onnx.data
    //               (ORT finds the .data file automatically; both must coexist)
    // tokenizerJson — path to tokenizer.json from the model directory
    EOUDetector(const std::string& modelPath,
                const std::string& tokenizerJson,
                bool useGpu = false)
        : env_(ORT_LOGGING_LEVEL_WARNING, "eou")
        , memoryInfo_("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault)
    {
        // ── session ──────────────────────────────────────────────────────────
        Ort::SessionOptions opts;
        opts.SetIntraOpNumThreads(1);   // per-request; tune for your CPU core budget
        opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
        if (useGpu) {
            OrtCUDAProviderOptions cuda{};
            cuda.device_id = 0;
            opts.AppendExecutionProvider_CUDA(cuda);
        }

        std::string onnxPath = modelPath + "/model.onnx";
        session_ = std::make_unique<Ort::Session>(env_, onnxPath.c_str(), opts);

        // ── cache I/O names (avoids per-call allocation) ─────────────────────
        inputName_  = session_->GetInputNameAllocated(0, allocator_).get();
        outputName_ = session_->GetOutputNameAllocated(0, allocator_).get();

        // ── tokenizer ────────────────────────────────────────────────────────
        std::ifstream f(tokenizerJson);
        std::string json((std::istreambuf_iterator<char>(f)),
                          std::istreambuf_iterator<char>());
        tokenizer_ = tokenizers::Tokenizer::FromBlobJSON(json);
    }

    // Returns P(end-of-utterance) in [0, 1].
    // history    — prior turns in chronological order, alternating user/assistant
    //              (may be empty for first turn)
    // currentText — the latest partial or complete utterance being evaluated
    float score(const std::vector<Turn>& history,
                const std::string&       currentText)
    {
        std::vector<int64_t> ids = buildInputIds(history, currentText);

        // ── dims ─────────────────────────────────────────────────────────────
        int64_t seqLen = static_cast<int64_t>(ids.size());
        std::vector<int64_t> inputDims  {1, seqLen};
        std::vector<int64_t> outputDims {1};

        // ── I/O name pointers ─────────────────────────────────────────────────
        std::vector<const char*> inputNames  {inputName_.c_str()};
        std::vector<const char*> outputNames {outputName_.c_str()};

        // ── input tensor (INT64) ──────────────────────────────────────────────
        std::vector<Ort::Value> inputTensors;
        inputTensors.push_back(
            Ort::Value::CreateTensor<int64_t>(
                memoryInfo_,
                ids.data(), ids.size(),
                inputDims.data(), inputDims.size()));

        // ── output tensor (FLOAT32) ───────────────────────────────────────────
        std::vector<float> outputValues(1);
        std::vector<Ort::Value> outputTensors;
        outputTensors.push_back(
            Ort::Value::CreateTensor<float>(
                memoryInfo_,
                outputValues.data(), outputValues.size(),
                outputDims.data(), outputDims.size()));

        // ── run ───────────────────────────────────────────────────────────────
        session_->Run(
            Ort::RunOptions{nullptr},
            inputNames.data(),  inputTensors.data(),  inputTensors.size(),
            outputNames.data(), outputTensors.data(), outputTensors.size());

        return outputValues[0];
    }

    bool isEndOfUtterance(const std::vector<Turn>& history,
                          const std::string&        currentText)
    {
        return score(history, currentText) >= THRESHOLD;
    }

private:
    // ── Qwen2.5 chat-template token IDs (hardcoded — never change) ───────────
    //
    // <|im_start|>system\n
    // You are Qwen, created by Alibaba Cloud. You are a helpful assistant.
    // <|im_end|>\n
    static const std::vector<int64_t> SYSTEM_BLOCK;
    static const std::vector<int64_t> USER_OPEN;    // <|im_start|>user\n
    static const std::vector<int64_t> USER_CLOSE;   // <|im_end|>\n
    static const std::vector<int64_t> ASST_OPEN;    // <|im_start|>assistant\n
    static const std::vector<int64_t> ASST_CLOSE;   // <|im_end|>\n

    // ── members ───────────────────────────────────────────────────────────────
    Ort::Env                              env_;
    Ort::AllocatorWithDefaultOptions      allocator_;
    Ort::MemoryInfo                       memoryInfo_;
    std::unique_ptr<Ort::Session>         session_;
    std::unique_ptr<tokenizers::Tokenizer> tokenizer_;
    std::string                           inputName_;
    std::string                           outputName_;

    // ── text preprocessing ────────────────────────────────────────────────────
    // Matches training: NFKC → lowercase → strip punctuation (keep ' -)
    // → collapse whitespace
    std::string preprocess(const std::string& raw)
    {
        // NFKC normalise + fold case via ICU
        icu::UnicodeString u = icu::UnicodeString::fromUTF8(raw);
        UErrorCode err = U_ZERO_ERROR;
        const icu::Normalizer2* nfkc =
            icu::Normalizer2::getNFKCInstance(err);
        u = nfkc->normalize(u, err);
        u.foldCase(U_FOLD_CASE_DEFAULT);

        std::string s;
        u.toUTF8String(s);

        // strip punctuation (keep apostrophe and hyphen), collapse spaces
        std::string out;
        out.reserve(s.size());
        bool lastWasSpace = true;
        for (unsigned char c : s) {
            if (c == '\'' || c == '-') {
                out += c; lastWasSpace = false;
            } else if (std::ispunct(c)) {
                // skip
            } else if (std::isspace(c)) {
                if (!lastWasSpace) { out += ' '; lastWasSpace = true; }
            } else {
                out += c; lastWasSpace = false;
            }
        }
        while (!out.empty() && out.back() == ' ') out.pop_back();
        return out;
    }

    // ── BPE-encode text, return int64 ids ────────────────────────────────────
    std::vector<int64_t> encode(const std::string& text)
    {
        auto ids32 = tokenizer_->Encode(text, /*add_special_tokens=*/false);
        std::vector<int64_t> ids64(ids32.begin(), ids32.end());
        return ids64;
    }

    // ── append helper ─────────────────────────────────────────────────────────
    static void append(std::vector<int64_t>& dst, const std::vector<int64_t>& src)
    {
        dst.insert(dst.end(), src.begin(), src.end());
    }

    // ── build full input_ids for ONNX ────────────────────────────────────────
    std::vector<int64_t> buildInputIds(const std::vector<Turn>& history,
                                       const std::string&        currentText)
    {
        std::vector<int64_t> ids;
        ids.reserve(MAX_TOKENS);

        // system block (always present)
        append(ids, SYSTEM_BLOCK);

        // prior turns — merge consecutive same-role turns
        std::vector<Turn> merged;
        for (const auto& t : history) {
            std::string clean = preprocess(t.text);
            if (!merged.empty() && merged.back().role == t.role)
                merged.back().text += " " + clean;
            else
                merged.push_back({t.role, clean});
        }

        for (const auto& t : merged) {
            if (t.role == Turn::Role::User) {
                append(ids, USER_OPEN);
                append(ids, encode(t.text));
                append(ids, USER_CLOSE);
            } else {
                append(ids, ASST_OPEN);
                append(ids, encode(t.text));
                append(ids, ASST_CLOSE);
            }
        }

        // current (last) user turn — NO closing <|im_end|>
        append(ids, USER_OPEN);
        append(ids, encode(preprocess(currentText)));

        // truncate to MAX_TOKENS from the right (keep most recent context)
        if ((int)ids.size() > MAX_TOKENS)
            ids = std::vector<int64_t>(ids.end() - MAX_TOKENS, ids.end());

        return ids;
    }
};

// ── static member definitions ─────────────────────────────────────────────────

const std::vector<int64_t> EOUDetector::SYSTEM_BLOCK = {
    151644, 8948, 198,                                    // <|im_start|>system\n
    2610, 525, 1207, 16948, 11, 3465, 553, 54364, 14817, // You are Qwen, created by Alibaba Cloud.
    13, 1446, 525, 264, 10950, 17847, 13,                 // You are a helpful assistant.
    151645, 198                                           // <|im_end|>\n
};
const std::vector<int64_t> EOUDetector::USER_OPEN  = {151644, 872,   198};
const std::vector<int64_t> EOUDetector::USER_CLOSE = {151645,        198};
const std::vector<int64_t> EOUDetector::ASST_OPEN  = {151644, 77091, 198};
const std::vector<int64_t> EOUDetector::ASST_CLOSE = {151645,        198};