// pinta_router.cpp
// AENEA Pinta-2.0 — 74M semantic router, edge-first.
// No calibration. No hardcoded paths. Drop folder anywhere, it works.
//
// Build (MSVC x64):
//   cl /std:c++17 /O2 /EHsc /I"path\to\ort\include" pinta_router.cpp
//      /link /LIBPATH:"path\to\ort\lib" onnxruntime.lib /out:pinta_router.exe
//
// Run:
//   pinta_router.exe "Design a microservices architecture"
//   pinta_router.exe --model M.onnx --tokenizer T.json "some prompt"

#include <onnxruntime_cxx_api.h>
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <map>
#include <numeric>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>

#ifdef _WIN32
#include <windows.h>
#else
#include <libgen.h>
#include <unistd.h>
#include <limits.h>
#endif

namespace {

// ── constants ────────────────────────────────────────────────────────────────
constexpr int64_t kNumClasses = 9;
constexpr int     kWarmup     = 5;
constexpr int     kBench      = 200;
constexpr float   kConfThresh = 0.35f;
constexpr int     kFallback   = 6;  // reserved_29 = Knowledge
constexpr int     kMaxSeq     = 2048;

const char* kDomainNames[9] = {
    "SyntaxDevOps", "Code", "Creative", "RAG", "Architecture",
    "Math", "Knowledge", "Compliance", "Chat"
};
const char* kDomainTokens[9] = {
    "<|reserved_23|>", "<|reserved_24|>", "<|reserved_25|>",
    "<|reserved_26|>", "<|reserved_27|>", "<|reserved_28|>",
    "<|reserved_29|>", "<|reserved_30|>", "<|reserved_31|>"
};

// ── portable exe directory ───────────────────────────────────────────────────
std::string exe_dir() {
#ifdef _WIN32
    char buf[MAX_PATH];
    GetModuleFileNameA(nullptr, buf, MAX_PATH);
    std::string p(buf);
    size_t pos = p.find_last_of("\\/");
    return (pos == std::string::npos) ? "." : p.substr(0, pos);
#else
    char buf[PATH_MAX];
    ssize_t n = readlink("/proc/self/exe", buf, sizeof(buf) - 1);
    if (n <= 0) return ".";
    buf[n] = '\0';
    std::string p(buf);
    size_t pos = p.find_last_of('/');
    return (pos == std::string::npos) ? "." : p.substr(0, pos);
#endif
}

std::string path_join(const std::string& dir, const std::string& file) {
#ifdef _WIN32
    return dir + "\\" + file;
#else
    return dir + "/" + file;
#endif
}

bool file_exists(const std::string& p) {
    std::ifstream f(p);
    return f.good();
}

// ── GPT-2 byte encoder ───────────────────────────────────────────────────────
struct ByteEnc { std::string s[256]; };

ByteEnc build_byte_enc() {
    ByteEnc be;
    std::vector<int> bs, cs;
    for (int b = '!'; b <= '~'; ++b) bs.push_back(b);
    for (int b = 0xA1; b <= 0xAC; ++b) bs.push_back(b);
    for (int b = 0xAE; b <= 0xFF; ++b) bs.push_back(b);
    cs = bs;
    int n = 0;
    for (int b = 0; b < 256; ++b) {
        if (std::find(bs.begin(), bs.end(), b) == bs.end()) {
            bs.push_back(b);
            cs.push_back(256 + n);
            ++n;
        }
    }
    for (size_t i = 0; i < bs.size(); ++i) {
        char32_t cp = static_cast<char32_t>(cs[i]);
        char buf[5] = {};
        if (cp < 0x80) {
            buf[0] = static_cast<char>(cp);
        } else if (cp < 0x800) {
            buf[0] = static_cast<char>(0xC0 | (cp >> 6));
            buf[1] = static_cast<char>(0x80 | (cp & 0x3F));
        } else if (cp < 0x10000) {
            buf[0] = static_cast<char>(0xE0 | (cp >> 12));
            buf[1] = static_cast<char>(0x80 | ((cp >> 6) & 0x3F));
            buf[2] = static_cast<char>(0x80 | (cp & 0x3F));
        } else {
            buf[0] = static_cast<char>(0xF0 | (cp >> 18));
            buf[1] = static_cast<char>(0x80 | ((cp >> 12) & 0x3F));
            buf[2] = static_cast<char>(0x80 | ((cp >> 6) & 0x3F));
            buf[3] = static_cast<char>(0x80 | (cp & 0x3F));
        }
        be.s[bs[i]] = std::string(buf);
    }
    return be;
}

// ── minimal JSON parser ──────────────────────────────────────────────────────
enum class JT { Null, Bool, Num, Str, Arr, Obj };

struct JV {
    JT t = JT::Null;
    std::string s;
    double num = 0;
    std::vector<JV> arr;
    std::vector<std::pair<std::string, JV>> obj;
    const JV* get(const std::string& k) const {
        if (t != JT::Obj) return nullptr;
        for (auto& kv : obj)
            if (kv.first == k) return &kv.second;
        return nullptr;
    }
};

class JsonParser {
public:
    explicit JsonParser(const std::string& src) : s_(src) {}
    JV parse() { skip_ws(); JV v = parse_value(); skip_ws(); return v; }

private:
    const std::string& s_;
    size_t p_ = 0;

    void die(const char* msg) {
        std::fprintf(stderr, "[json] error at pos %zu: %s\n", p_, msg);
        std::exit(1);
    }
    char peek() const { return p_ < s_.size() ? s_[p_] : '\0'; }
    char take() { return p_ < s_.size() ? s_[p_++] : '\0'; }
    void skip_ws() {
        while (p_ < s_.size() && (s_[p_] == ' ' || s_[p_] == '\t' || s_[p_] == '\n' || s_[p_] == '\r'))
            ++p_;
    }
    void expect(char c) { skip_ws(); if (take() != c) die("unexpected char"); }

    JV parse_value() {
        skip_ws();
        char c = peek();
        if (c == '"') return parse_string();
        if (c == '{') return parse_object();
        if (c == '[') return parse_array();
        if (c == 't') { p_ += 4; JV v; v.t = JT::Bool; v.s = "true"; return v; }
        if (c == 'f') { p_ += 5; JV v; v.t = JT::Bool; v.s = "false"; return v; }
        if (c == 'n') { p_ += 4; return JV{}; }
        return parse_number();
    }

    JV parse_string() {
        expect('"');
        JV v; v.t = JT::Str;
        while (true) {
            char c = take();
            if (c == '\0') die("unterminated string");
            if (c == '"') break;
            if (c == '\\') {
                char e = take();
                switch (e) {
                    case '"':  v.s += '"';  break;
                    case '\\': v.s += '\\'; break;
                    case '/':  v.s += '/';  break;
                    case 'b':  v.s += '\b'; break;
                    case 'f':  v.s += '\f'; break;
                    case 'n':  v.s += '\n'; break;
                    case 'r':  v.s += '\r'; break;
                    case 't':  v.s += '\t'; break;
                    case 'u':  for (int i = 0; i < 4; ++i) take(); v.s += '?'; break;
                    default:   die("bad escape");
                }
            } else {
                v.s += c;
            }
        }
        return v;
    }

    JV parse_object() {
        expect('{');
        JV v; v.t = JT::Obj;
        skip_ws();
        if (peek() == '}') { take(); return v; }
        while (true) {
            JV key = parse_string();
            expect(':');
            JV val = parse_value();
            v.obj.emplace_back(std::move(key.s), std::move(val));
            skip_ws();
            char c = take();
            if (c == '}') break;
            if (c != ',') die("expected , or }");
        }
        return v;
    }

    JV parse_array() {
        expect('[');
        JV v; v.t = JT::Arr;
        skip_ws();
        if (peek() == ']') { take(); return v; }
        while (true) {
            v.arr.push_back(parse_value());
            skip_ws();
            char c = take();
            if (c == ']') break;
            if (c != ',') die("expected , or ]");
        }
        return v;
    }

    JV parse_number() {
        size_t start = p_;
        if (peek() == '-') ++p_;
        while (p_ < s_.size() && (isdigit((unsigned char)s_[p_]) || s_[p_] == '.' ||
               s_[p_] == 'e' || s_[p_] == 'E' || s_[p_] == '+' || s_[p_] == '-'))
            ++p_;
        JV v; v.t = JT::Num;
        v.s = s_.substr(start, p_ - start);
        v.num = v.s.empty() ? 0.0 : std::stod(v.s);
        return v;
    }
};

// ── prompt template ──────────────────────────────────────────────────────────
std::string wrap_prompt(const std::string& text) {
    return "User: " + text + "\nAssistant: ";
}

// ── tokenizer ────────────────────────────────────────────────────────────────
class Tokenizer {
public:
    explicit Tokenizer(const std::string& path) : byte_(build_byte_enc()) {
        std::ifstream f(path, std::ios::binary);
        if (!f) { std::fprintf(stderr, "[tok] cannot open: %s\n", path.c_str()); std::exit(1); }
        std::stringstream ss;
        ss << f.rdbuf();
        std::string raw = ss.str();

        JsonParser jp(raw);
        JV root = jp.parse();

        // vocab + merges
        const JV* model = root.get("model");
        if (model) {
            if (const JV* v = model->get("vocab"))
                for (auto& kv : v->obj)
                    vocab_[kv.first] = static_cast<int64_t>(kv.second.num);
            if (const JV* m = model->get("merges")) {
                for (size_t i = 0; i < m->arr.size(); ++i) {
                    const JV& e = m->arr[i];
                    std::string a, b;
                    if (e.t == JT::Arr && e.arr.size() >= 2) {
                        a = e.arr[0].s; b = e.arr[1].s;
                    } else if (e.t == JT::Str) {
                        size_t sp = e.s.find(' ');
                        a = e.s.substr(0, sp);
                        b = (sp == std::string::npos) ? "" : e.s.substr(sp + 1);
                    } else continue;
                    ranks_[std::make_pair(a, b)] = static_cast<int64_t>(i);
                }
            }
        }

        // added tokens (reserved_23..31 live here)
        if (const JV* at = root.get("added_tokens")) {
            for (auto& e : at->arr) {
                const JV* c = e.get("content");
                const JV* id = e.get("id");
                if (!c || !id) continue;
                added_[c->s] = static_cast<int64_t>(id->num);
                specials_.push_back(c->s);
            }
        }
        std::sort(specials_.begin(), specials_.end(),
                  [](const std::string& a, const std::string& b) { return a.size() > b.size(); });

        // pad / unk
        pad_ = lookup("<|padding|>"); if (pad_ < 0) pad_ = lookup("<pad>"); if (pad_ < 0) pad_ = 0;
        unk_ = lookup("<|unk|>");     if (unk_ < 0) unk_ = lookup("<unk>"); if (unk_ < 0) unk_ = pad_;

        // routing IDs
        for (int i = 23; i <= 31; ++i) {
            std::string t = "<|reserved_" + std::to_string(i) + "|>";
            int64_t id = lookup(t);
            if (id < 0) {
                id = i;
                std::fprintf(stderr, "[tok] warn: %s not in vocab, fallback id=%d\n", t.c_str(), i);
            }
            routing_ids_.push_back(id);
        }
    }

    int64_t pad_id() const { return pad_; }
    const std::vector<int64_t>& routing_ids() const { return routing_ids_; }

    std::vector<int64_t> encode(const std::string& text) const {
        std::vector<int64_t> ids;
        size_t pos = 0;
        while (pos < text.size()) {
            size_t hit = std::string::npos;
            const std::string* which = nullptr;
            for (auto& sp : specials_) {
                size_t f = text.find(sp, pos);
                if (f != std::string::npos && (hit == std::string::npos || f < hit)) {
                    hit = f; which = &sp;
                }
            }
            size_t end = (hit == std::string::npos) ? text.size() : hit;
            if (end > pos) encode_text(text.substr(pos, end - pos), ids);
            if (which) {
                auto a = added_.find(*which);
                ids.push_back(a != added_.end() ? a->second : unk_);
                pos = hit + which->size();
            } else {
                pos = end;
            }
        }
        return ids;
    }

private:
    int64_t lookup(const std::string& t) const {
        auto a = vocab_.find(t); if (a != vocab_.end()) return a->second;
        auto b = added_.find(t); if (b != added_.end()) return b->second;
        return -1;
    }

    static bool is_letter(unsigned char c) { return c >= 0x80 || std::isalpha(c); }
    static bool is_digit(unsigned char c)  { return std::isdigit(c) != 0; }
    static unsigned char char_class(unsigned char c) {
        if (is_letter(c)) return 1;
        if (is_digit(c))  return 2;
        return 3;
    }

    void encode_text(const std::string& chunk, std::vector<int64_t>& ids) const {
        size_t i = 0, n = chunk.size();
        while (i < n) {
            std::string piece;
            if (chunk[i] == ' ') {
                size_t j = i;
                while (j < n && chunk[j] == ' ') ++j;
                size_t spaces = j - i;
                if (spaces == 1 && j < n) {
                    piece += byte_.s[(unsigned char)chunk[j]];
                    unsigned char cls = char_class((unsigned char)chunk[j]);
                    ++j;
                    while (j < n && char_class((unsigned char)chunk[j]) == cls && chunk[j] != ' ') {
                        piece += byte_.s[(unsigned char)chunk[j]];
                        ++j;
                    }
                } else {
                    for (size_t k = i; k < j; ++k)
                        piece += byte_.s[(unsigned char)chunk[k]];
                }
                i = j;
            } else {
                unsigned char cls = char_class((unsigned char)chunk[i]);
                while (i < n && chunk[i] != ' ' && char_class((unsigned char)chunk[i]) == cls) {
                    piece += byte_.s[(unsigned char)chunk[i]];
                    ++i;
                }
            }
            if (!piece.empty()) bpe(piece, ids);
        }
    }

    void bpe(const std::string& piece, std::vector<int64_t>& ids) const {
        std::vector<std::string> word;
        for (unsigned char b : piece) word.push_back(byte_.s[b]);
        while (word.size() >= 2) {
            int64_t best_rank = -1;
            size_t best_at = 0;
            bool found = false;
            for (size_t k = 0; k + 1 < word.size(); ++k) {
                auto it = ranks_.find(std::make_pair(word[k], word[k + 1]));
                if (it != ranks_.end() && (!found || it->second < best_rank)) {
                    best_rank = it->second;
                    best_at = k;
                    found = true;
                }
            }
            if (!found) break;
            std::string merged = word[best_at] + word[best_at + 1];
            word.erase(word.begin() + best_at, word.begin() + best_at + 2);
            word.insert(word.begin() + best_at, merged);
        }
        for (auto& sym : word) {
            auto it = vocab_.find(sym);
            ids.push_back(it != vocab_.end() ? it->second : unk_);
        }
    }

    ByteEnc byte_;
    std::unordered_map<std::string, int64_t> vocab_, added_;
    std::map<std::pair<std::string, std::string>, int64_t> ranks_;
    std::vector<std::string> specials_;
    std::vector<int64_t> routing_ids_;
    int64_t pad_ = 0, unk_ = 0;
};

// ── ONNX router ──────────────────────────────────────────────────────────────
class Router {
public:
    Router(const std::string& model_path, const std::string& tok_path) : tok_(tok_path) {
        env_ = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "pinta");
        Ort::SessionOptions so;
        so.SetIntraOpNumThreads(0);  // 0 = let ORT pick optimal thread count
        so.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
        so.EnableMemPattern();

        std::wstring wpath(model_path.begin(), model_path.end());
        session_ = Ort::Session(env_, wpath.c_str(), so);

        // detect input name
        Ort::AllocatorWithDefaultOptions alloc;
        size_t nin = session_.GetInputCount();
        input_name_ = session_.GetInputNameAllocated(0, alloc).get();
        for (size_t i = 0; i < nin; ++i) {
            std::string nm(session_.GetInputNameAllocated(i, alloc).get());
            if (nm.find("input_ids") != std::string::npos) { input_name_ = nm; break; }
        }

        // detect output name
        size_t nout = session_.GetOutputCount();
        output_name_ = session_.GetOutputNameAllocated(0, alloc).get();
        for (size_t i = 0; i < nout; ++i) {
            std::string nm(session_.GetOutputNameAllocated(i, alloc).get());
            if (nm.find("logit") != std::string::npos) { output_name_ = nm; break; }
        }

        in_ptrs_  = { input_name_.c_str() };
        out_ptrs_ = { output_name_.c_str() };
    }

    struct Result {
        int cls; float conf; double us;
        int raw_cls; float raw_conf;
        int top3[3]; float top3_p[3];
        bool fallback;
    };

    void prepare(const std::string& prompt) {
        cached_ids_ = tok_.encode(prompt);
        if (cached_ids_.empty()) cached_ids_.push_back(tok_.pad_id());
        if (cached_ids_.size() > kMaxSeq)
            cached_ids_.erase(cached_ids_.begin(), cached_ids_.end() - kMaxSeq);
        n_ = cached_ids_.size();
    }

    Result run() {
        // create tensor at actual sequence length — no padding
        int64_t shape[2] = { 1, static_cast<int64_t>(n_) };
        auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
        std::vector<Ort::Value> inputs;
        inputs.emplace_back(Ort::Value::CreateTensor<int64_t>(
            mem, cached_ids_.data(), cached_ids_.size(), shape, 2));

        auto t0 = std::chrono::high_resolution_clock::now();
        auto outs = session_.Run(run_opts_, in_ptrs_.data(), inputs.data(),
                                 inputs.size(), out_ptrs_.data(), 1);
        auto t1 = std::chrono::high_resolution_clock::now();

        auto& out = outs[0];
        auto oshape = out.GetTensorTypeAndShapeInfo().GetShape();
        const float* p = out.GetTensorData<float>();
        const auto& rid = tok_.routing_ids();

        float rl[kNumClasses];
        if (oshape.size() == 3) {
            int64_t V = oshape[2];
            int64_t row = std::min<int64_t>(static_cast<int64_t>(n_), oshape[1]) - 1;
            for (int64_t i = 0; i < kNumClasses; ++i) rl[i] = p[row * V + rid[i]];
        } else if (oshape.size() == 2) {
            int64_t d = oshape[1];
            if (d == kNumClasses) {
                for (int64_t i = 0; i < d; ++i) rl[i] = p[i];
            } else {
                for (int64_t i = 0; i < kNumClasses; ++i) rl[i] = p[rid[i]];
            }
        } else {
            std::fprintf(stderr, "[router] unsupported logits rank %zu\n", oshape.size());
            std::exit(1);
        }

        // softmax
        float mx = *std::max_element(rl, rl + kNumClasses);
        float sum = 0.f;
        float pr[kNumClasses];
        for (int i = 0; i < kNumClasses; ++i) { pr[i] = std::exp(rl[i] - mx); sum += pr[i]; }
        for (int i = 0; i < kNumClasses; ++i) pr[i] /= sum;

        int best = 0; float bp = -1.f;
        for (int i = 0; i < kNumClasses; ++i)
            if (pr[i] > bp) { bp = pr[i]; best = i; }

        bool fb = bp < kConfThresh;
        int final_cls = fb ? kFallback : best;

        int idx[9];
        for (int i = 0; i < 9; ++i) idx[i] = i;
        std::partial_sort(idx, idx + 3, idx + 9, [&](int a, int b) { return pr[a] > pr[b]; });

        double us = std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count();

        Result r;
        r.cls = final_cls;
        r.conf = fb ? pr[kFallback] : bp;
        r.raw_cls = best;
        r.raw_conf = bp;
        for (int i = 0; i < 3; ++i) { r.top3[i] = idx[i]; r.top3_p[i] = pr[idx[i]]; }
        r.fallback = fb;
        r.us = us;
        return r;
    }

private:
    Ort::Env env_{nullptr};
    Ort::Session session_{nullptr};
    Ort::RunOptions run_opts_;
    Tokenizer tok_;
    std::string input_name_, output_name_;
    std::vector<const char*> in_ptrs_, out_ptrs_;
    std::vector<int64_t> cached_ids_;
    size_t n_ = 0;
};

} // namespace

// ── main ─────────────────────────────────────────────────────────────────────
int main(int argc, char** argv) {
    std::fprintf(stderr,
        "===========================================================\n"
        " AENEA Pinta-2.0\n"
        " 74M semantic router | edge dispatch | Apache 2.0\n"
        "===========================================================\n\n");

    // resolve paths: same directory as exe, or CLI override
    std::string dir = exe_dir();
    std::string model_path = path_join(dir, "pinta-2.0.opset18.onnx");
    std::string tok_path   = path_join(dir, "tokenizer.json");
    std::string prompt_text = "Design a microservices architecture for our payment system.";

    // parse optional CLI args
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "--model" && i + 1 < argc) { model_path = argv[++i]; }
        else if (arg == "--tokenizer" && i + 1 < argc) { tok_path = argv[++i]; }
        else if (arg == "--help" || arg == "-h") {
            std::fprintf(stderr,
                "usage: pinta_router [prompt]\n"
                "       pinta_router --model M.onnx --tokenizer T.json [prompt]\n"
                "\n"
                "By default loads pinta-2.0.opset18.onnx and tokenizer.json\n"
                "from the same directory as the executable.\n");
            return 0;
        }
        else { prompt_text = arg; }
    }

    if (!file_exists(model_path)) {
        std::fprintf(stderr, "[error] model not found: %s\n", model_path.c_str());
        std::fprintf(stderr, "        place pinta-2.0.opset18.onnx next to the exe\n");
        return 1;
    }
    if (!file_exists(tok_path)) {
        std::fprintf(stderr, "[error] tokenizer not found: %s\n", tok_path.c_str());
        std::fprintf(stderr, "        place tokenizer.json next to the exe\n");
        return 1;
    }

    std::string prompt = wrap_prompt(prompt_text);
    Router router(model_path, tok_path);
    router.prepare(prompt);

    // warmup
    for (int i = 0; i < kWarmup; ++i) router.run();

    // single result
    Router::Result r = router.run();

    std::printf("model:      %s\n", model_path.c_str());
    std::printf("prompt:     \"%s\"\n", prompt_text.c_str());
    std::printf("token:      %s\n", kDomainTokens[r.cls]);
    std::printf("domain:     %s\n", kDomainNames[r.cls]);
    std::printf("raw argmax: %s  conf=%.4f\n", kDomainTokens[r.raw_cls], r.raw_conf);
    std::printf("fallback:   %s\n", r.fallback ? "YES -> Knowledge" : "no");
    std::printf("top-3:      %s %.3f | %s %.3f | %s %.3f\n",
        kDomainTokens[r.top3[0]], r.top3_p[0],
        kDomainTokens[r.top3[1]], r.top3_p[1],
        kDomainTokens[r.top3[2]], r.top3_p[2]);
    std::printf("latency:    %.1f us (%.3f ms)\n", r.us, r.us / 1000.0);

    // benchmark
    std::vector<double> times;
    times.reserve(kBench);
    for (int i = 0; i < kBench; ++i) times.push_back(router.run().us);
    std::sort(times.begin(), times.end());
    double mean = std::accumulate(times.begin(), times.end(), 0.0) / times.size();

    std::printf("\nbenchmark (%d warm runs):\n", kBench);
    std::printf("  min:  %8.1f us (%6.3f ms)\n", times.front(), times.front() / 1000.0);
    std::printf("  p50:  %8.1f us (%6.3f ms)\n", times[times.size() / 2], times[times.size() / 2] / 1000.0);
    std::printf("  p95:  %8.1f us (%6.3f ms)\n", times[(size_t)(0.95 * times.size())], times[(size_t)(0.95 * times.size())] / 1000.0);
    std::printf("  mean: %8.1f us (%6.3f ms)\n", mean, mean / 1000.0);
    std::printf("  max:  %8.1f us (%6.3f ms)\n", times.back(), times.back() / 1000.0);

    return 0;
}