// Prototype E2E OuteTTS-1.0 : speaker JSON (v3) + texte -> GGUF talker (libllama) -> codes DAC -> ONNX decoder (ORT) -> WAV 24 kHz
// Usage: outetts_runner <talker.gguf> <speaker.json> <dac.onnx> "<texte>" <out.wav> [n_threads=6] [ort_threads=2] [seed=42]
#include "llama.h"
#include <onnxruntime_cxx_api.h>
#include <nlohmann/json.hpp>

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <deque>
#include <fstream>
#include <random>
#include <sstream>
#include <string>
#include <vector>

using json = nlohmann::json;
static double now_s() {
    return std::chrono::duration<double>(std::chrono::steady_clock::now().time_since_epoch()).count();
}

// ---- prompt v3 (port de outetts/version/v3/prompt_processor.py, sans ftfy/NFKC : textes FR propres attendus)
static std::string build_prompt(const json & spk, const std::string & input_text) {
    std::string text = spk["text"].get<std::string>();
    // merge_speaker_text : separateur ". " si la phrase speaker ne finit pas par ./?/!
    std::string sep;
    while (!text.empty() && text.back() == ' ') text.pop_back();
    char last = text.empty() ? '.' : text.back();
    if (last != '.' && last != '?' && last != '!') sep = ".";
    std::string merged = text + sep + " " + input_text;

    std::ostringstream p;
    p << "<|im_start|>\n<|text_start|>" << merged << "<|text_end|>\n<|audio_start|>\n";

    const auto & words = spk["words"];
    for (size_t i = 0; i < words.size(); i++) {
        const auto & w = words[i];
        std::string wtxt = w["word"].get<std::string>();
        if (i + 1 == words.size()) wtxt += sep; // le dernier mot speaker porte le separateur
        char tbuf[32];
        snprintf(tbuf, sizeof(tbuf), "<|t_%.2f|>", w["duration"].get<double>());
        p << "<|word_start|>" << wtxt << "<|features|>" << tbuf
          << "<|energy_" << w["features"]["energy"].get<int>() << "|>"
          << "<|spectral_centroid_" << w["features"]["spectral_centroid"].get<int>() << "|>"
          << "<|pitch_" << w["features"]["pitch"].get<int>() << "|>"
          << "<|code|>";
        const auto & c1 = w["c1"]; const auto & c2 = w["c2"];
        for (size_t k = 0; k < c1.size(); k++)
            p << "<|c1_" << c1[k].get<int>() << "|><|c2_" << c2[k].get<int>() << "|>";
        p << "<|word_end|>";
        if (i + 1 < words.size()) p << "\n";
    }
    p << "\n<|word_start|>";
    return p.str();
}

static llama_token tok_id(const llama_vocab * vocab, const char * s) {
    llama_token t[8];
    int n = llama_tokenize(vocab, s, (int)strlen(s), t, 8, false, true);
    if (n != 1) { fprintf(stderr, "token special inattendu '%s' -> %d tokens\n", s, n); exit(1); }
    return t[0];
}

static void write_wav(const char * path, const float * x, size_t n, int sr) {
    std::ofstream f(path, std::ios::binary);
    auto w32 = [&](uint32_t v) { f.write((char *)&v, 4); };
    auto w16 = [&](uint16_t v) { f.write((char *)&v, 2); };
    f.write("RIFF", 4); w32(36 + (uint32_t)n * 2); f.write("WAVE", 4);
    f.write("fmt ", 4); w32(16); w16(1); w16(1); w32(sr); w32(sr * 2); w16(2); w16(16);
    f.write("data", 4); w32((uint32_t)n * 2);
    for (size_t i = 0; i < n; i++) {
        float v = x[i] < -1.f ? -1.f : (x[i] > 1.f ? 1.f : x[i]);
        int16_t s = (int16_t)(v * 32767.f);
        f.write((char *)&s, 2);
    }
}

int main(int argc, char ** argv) {
    if (argc < 6) { fprintf(stderr, "usage: %s talker.gguf speaker.json dac.onnx \"texte\" out.wav [t] [ort_t] [seed]\n", argv[0]); return 1; }
    const char * gguf = argv[1], * spk_path = argv[2], * dac_path = argv[3], * text = argv[4], * out_path = argv[5];
    int n_threads = argc > 6 ? atoi(argv[6]) : 6;
    int ort_threads = argc > 7 ? atoi(argv[7]) : 2;
    uint32_t seed = argc > 8 ? (uint32_t)atoi(argv[8]) : 42;
    bool restricted = argc > 9 ? atoi(argv[9]) != 0 : true;  // sampling restreint a l'alphabet utile
    bool kv8 = argc > 10 ? atoi(argv[10]) != 0 : false;      // KV cache q8_0 + flash-attn

    json spk; { std::ifstream f(spk_path); f >> spk; }
    std::string prompt = build_prompt(spk, text);

    llama_backend_init();
    llama_model_params mp = llama_model_default_params();
    mp.n_gpu_layers = 0;
    static ggml_backend_dev_t no_devices[] = { nullptr };
    mp.devices = no_devices; // equivalent -dev none : pas d'enumeration HTP residuelle
    llama_model * model = llama_model_load_from_file(gguf, mp);
    if (!model) { fprintf(stderr, "load fail %s\n", gguf); return 1; }
    const llama_vocab * vocab = llama_model_get_vocab(model);

    llama_context_params cp = llama_context_default_params();
    cp.n_ctx = argc > 11 ? atoi(argv[11]) : 8192; cp.n_batch = 512; cp.n_ubatch = 512;
    cp.n_threads = n_threads; cp.n_threads_batch = n_threads;
    if (kv8) {
        cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; // requis pour V quantise
        cp.type_k = GGML_TYPE_Q8_0;
        cp.type_v = GGML_TYPE_Q8_0;
    }
    llama_context * ctx = llama_init_from_model(model, cp);

    // ids de reference (bases contigues verifiees: c1=[<|c1_0|>..], c2 idem)
    const llama_token id_c1_0 = tok_id(vocab, "<|c1_0|>");
    const llama_token id_c2_0 = tok_id(vocab, "<|c2_0|>");
    const llama_token id_audio_end = tok_id(vocab, "<|audio_end|>");
    const llama_token id_im_end = tok_id(vocab, "<|im_end|>");

    // tokenize prompt
    std::vector<llama_token> ptok(prompt.size() + 16);
    int n_prompt = llama_tokenize(vocab, prompt.c_str(), (int)prompt.size(), ptok.data(), (int)ptok.size(), false, true);
    if (n_prompt < 0) { fprintf(stderr, "tokenize fail\n"); return 1; }
    ptok.resize(n_prompt);
    fprintf(stderr, "prompt: %d tokens\n", n_prompt);

    // sampler = defaults outetts: rep penalty 1.1 (fenetre 64), temp 0.4, top_k 40, top_p 0.9 (ordre HF)
    // mode restreint : on ne considere que l'alphabet generable -> evite le tri de 157k logits/token
    std::vector<llama_token> cand;
    if (restricted) {
        const llama_token id_specials_lo = tok_id(vocab, "<|endoftext|>"); // 151643, debut zone speciale
        const int n_vocab = llama_vocab_n_tokens(vocab);
        for (llama_token t = id_specials_lo; t < n_vocab; t++) cand.push_back(t); // specials + c1/c2 + features
        // tokens du texte cible (les mots sont regeneres dans les blocs) : variantes avec/sans espace
        std::string variants = std::string(text) + " " + text;
        std::vector<llama_token> ttok(variants.size() + 16);
        int nt = llama_tokenize(vocab, variants.c_str(), (int)variants.size(), ttok.data(), (int)ttok.size(), false, false);
        for (int i = 0; i < nt; i++) cand.push_back(ttok[i]);
        llama_token nl[4];
        int nn = llama_tokenize(vocab, "\n", 1, nl, 4, false, false);
        for (int i = 0; i < nn; i++) cand.push_back(nl[i]);
        std::sort(cand.begin(), cand.end());
        cand.erase(std::unique(cand.begin(), cand.end()), cand.end());
        fprintf(stderr, "alphabet restreint: %zu / %d tokens\n", cand.size(), n_vocab);
    }
    std::mt19937 rng(seed);
    std::deque<llama_token> last64;
    struct Cd { llama_token id; float logit; };
    auto sample_next = [&](void) -> llama_token {
        const float * logits = llama_get_logits_ith(ctx, -1);
        static std::vector<Cd> cds;
        cds.clear();
        if (restricted) {
            for (llama_token t : cand) cds.push_back({t, logits[t]});
        } else {
            const int n_vocab = llama_vocab_n_tokens(vocab);
            for (llama_token t = 0; t < n_vocab; t++) cds.push_back({t, logits[t]});
        }
        // rep penalty HF sur la fenetre 64
        for (auto & c : cds) {
            for (llama_token p : last64) {
                if (p == c.id) { c.logit = c.logit > 0 ? c.logit / 1.1f : c.logit * 1.1f; break; }
            }
        }
        // temp -> top_k -> top_p (ordre HF)
        for (auto & c : cds) c.logit /= 0.4f;
        size_t k = cds.size() < 40 ? cds.size() : 40;
        std::partial_sort(cds.begin(), cds.begin() + k, cds.end(), [](const Cd & a, const Cd & b) { return a.logit > b.logit; });
        cds.resize(k);
        float mx = cds[0].logit, sum = 0;
        for (auto & c : cds) { c.logit = expf(c.logit - mx); sum += c.logit; }
        float cum = 0; size_t keep = k;
        for (size_t i = 0; i < k; i++) { cum += cds[i].logit / sum; if (cum >= 0.9f) { keep = i + 1; break; } }
        cds.resize(keep);
        float sum2 = 0; for (auto & c : cds) sum2 += c.logit;
        std::uniform_real_distribution<float> U(0.f, sum2);
        float r = U(rng), acc = 0;
        llama_token pick = cds.back().id;
        for (auto & c : cds) { acc += c.logit; if (acc >= r) { pick = c.id; break; } }
        last64.push_back(pick);
        if (last64.size() > 64) last64.pop_front();
        return pick;
    };

    // session DAC creee AVANT (en prod: une fois au demarrage)
    double t0 = now_s();
    Ort::Env env(ORT_LOGGING_LEVEL_ERROR, "dac");
    Ort::SessionOptions so;
    so.SetIntraOpNumThreads(ort_threads);
    so.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
    Ort::Session session(env, dac_path, so);
    fprintf(stderr, "dac session init: %.2fs\n", now_s() - t0);

    // prefill par chunks de n_batch
    t0 = now_s();
    for (int off = 0; off < n_prompt; off += 512) {
        int n = n_prompt - off < 512 ? n_prompt - off : 512;
        llama_batch b = llama_batch_get_one(ptok.data() + off, n);
        if (llama_decode(ctx, b)) { fprintf(stderr, "prefill fail @%d\n", off); return 1; }
    }
    double t_prefill = now_s() - t0;
    llama_batch batch;

    // decode
    std::vector<int64_t> c1, c2;
    int n_gen = 0; const int n_max = 6000;
    t0 = now_s();
    llama_token cur = 0;
    double t_win = now_s();
    while (n_gen < n_max) {
        cur = sample_next();
        n_gen++;
        if (n_gen % 128 == 0) {
            double tn = now_s();
            fprintf(stderr, "  [%d] fenetre: %.1f tok/s (depth %d)\n", n_gen, 128.0 / (tn - t_win), n_prompt + n_gen);
            t_win = tn;
        }
        if (cur == id_audio_end || cur == id_im_end) break;
        if (cur >= id_c1_0 && cur < id_c1_0 + 1025) c1.push_back(cur - id_c1_0);
        else if (cur >= id_c2_0 && cur < id_c2_0 + 1025) c2.push_back(cur - id_c2_0);
        batch = llama_batch_get_one(&cur, 1);
        if (llama_decode(ctx, batch)) { fprintf(stderr, "decode fail @%d\n", n_gen); return 1; }
    }
    double t_gen = now_s() - t0;
    size_t n_frames = c1.size() < c2.size() ? c1.size() : c2.size();
    fprintf(stderr, "gen: %d tokens en %.2fs (%.1f tok/s) -> %zu frames (%.2fs audio)\n",
            n_gen, t_gen, n_gen / t_gen, n_frames, n_frames / 75.0);
    if (n_frames == 0) { fprintf(stderr, "aucun code genere\n"); return 1; }

    // DAC ONNX
    t0 = now_s();
    std::vector<int64_t> codes(2 * n_frames);
    memcpy(codes.data(), c1.data(), n_frames * sizeof(int64_t));
    memcpy(codes.data() + n_frames, c2.data(), n_frames * sizeof(int64_t));
    std::vector<int64_t> shape = {1, 2, (int64_t)n_frames};
    Ort::MemoryInfo mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value in = Ort::Value::CreateTensor<int64_t>(mem, codes.data(), codes.size(), shape.data(), 3);
    const char * in_names[] = {"codes"}; const char * out_names[] = {"audio"};
    auto out = session.Run(Ort::RunOptions{}, in_names, &in, 1, out_names, 1);
    double t_dac = now_s() - t0;
    auto oshape = out[0].GetTensorTypeAndShapeInfo().GetShape();
    size_t n_samples = (size_t)oshape[2];
    write_wav(out_path, out[0].GetTensorData<float>(), n_samples, 24000);

    double audio_s = n_samples / 24000.0;
    fprintf(stderr, "dac: %.2fs | audio %.2fs | prefill %.2fs\n", t_dac, audio_s, t_prefill);
    fprintf(stderr, "RTF gen=%.2f dac=%.2f total(sans prefill)=%.2f total(avec)=%.2f\n",
            t_gen / audio_s, t_dac / audio_s, (t_gen + t_dac) / audio_s, (t_prefill + t_gen + t_dac) / audio_s);
    printf("%s\n", out_path);

    llama_free(ctx);
    llama_model_free(model);
    llama_backend_free();
    return 0;
}