#include "qwen3asr.h"

#include <chrono>
#include <cstdio>
#include <vector>

#include <llama.h>
#include <mtmd.h>
#include <mtmd-helper.h>

namespace lc {

namespace {

std::string defaultQwenChatPrompt(const std::string & content)
{
    // Qwen3-ASR 的 chat template 就是标准 ChatML（见模型 GGUF 的 tokenizer.chat_template）
    return "<|im_start|>user\n" + content + "<|im_end|>\n<|im_start|>assistant\n";
}

void silentLog(ggml_log_level /*level*/, const char * /*text*/, void * /*user*/) {}

} // namespace

struct Qwen3Asr::Impl {
    AsrOptions        opt;
    llama_model *     model = nullptr;
    llama_context *   ctx   = nullptr;
    mtmd_context *    mctx  = nullptr;
    const llama_vocab * vocab = nullptr;
    int               audioRate = 16000;

    // 流式（增量）状态
    bool                      streamReady = false;
    std::vector<llama_token>  headTokens;   // 音频之前的 prompt 文本
    std::vector<llama_token>  tailTokens;   // 音频之后的 prompt 文本（含 assistant 头）
    std::vector<float>        streamAudio;  // 当前这段的全部音频（不滑动）
    size_t                    stableSamples = 0;   // 已固化进 KV 的音频样本数（8s 对齐）
    llama_pos                 nPastStable   = 0;   // stable 音频末尾的 KV 位置

    std::string buildPrompt(const std::string & marker) const
    {
        const char * tmpl = llama_model_chat_template(model, nullptr);
        if (tmpl) {
            llama_chat_message msg{"user", marker.c_str()};
            std::vector<char> buf(marker.size() * 2 + 4096);
            const int32_t n = llama_chat_apply_template(tmpl, &msg, 1, true, buf.data(), (int32_t) buf.size());
            if (n > 0 && (size_t) n <= buf.size()) {
                return std::string(buf.data(), (size_t) n);
            }
        }
        return defaultQwenChatPrompt(marker);
    }

    // 用显式位置解码一批 token（便于在增量流式里精确控制位置）
    bool decodeTokens(const llama_token * tokens, int32_t n, llama_pos startPos, bool wantLogits)
    {
        if (n <= 0) return true;
        llama_batch b = llama_batch_init(n, 0, 1);
        for (int32_t i = 0; i < n; ++i) {
            b.token[i]     = tokens[i];
            b.pos[i]       = startPos + i;
            b.n_seq_id[i]  = 1;
            b.seq_id[i][0] = 0;
            b.logits[i]    = (wantLogits && i == n - 1) ? 1 : 0;
        }
        b.n_tokens = n;
        const int32_t rc = llama_decode(ctx, b);
        llama_batch_free(b);
        return rc == 0;
    }

    // 把一段音频编码后追加到 KV 的 nPast 处
    bool appendAudio(const float * pcm, size_t nSamples, llama_pos & nPast, std::string * err)
    {
        if (nSamples == 0) return true;

        mtmd_bitmap * bmp = mtmd_bitmap_init_from_audio(nSamples, pcm);
        if (!bmp) { if (err) *err = "构造音频 bitmap 失败"; return false; }

        mtmd_input_chunks * chunks = mtmd_input_chunks_init();
        const std::string marker = mtmd_default_marker();
        const mtmd_input_text itext{ marker.c_str(), false, true };
        const mtmd_bitmap * bms[1] = { bmp };

        if (mtmd_tokenize(mctx, chunks, &itext, bms, 1) != 0) {
            mtmd_bitmap_free(bmp);
            mtmd_input_chunks_free(chunks);
            if (err) *err = "音频 mtmd_tokenize 失败";
            return false;
        }

        const mtmd_input_chunk * audioChunk = nullptr;
        for (size_t i = 0; i < mtmd_input_chunks_size(chunks); ++i) {
            const mtmd_input_chunk * c = mtmd_input_chunks_get(chunks, i);
            if (mtmd_input_chunk_get_type(c) == MTMD_INPUT_CHUNK_TYPE_AUDIO) { audioChunk = c; break; }
        }
        if (!audioChunk) {
            mtmd_bitmap_free(bmp);
            mtmd_input_chunks_free(chunks);
            if (err) *err = "tokenize 结果里没有音频块";
            return false;
        }

        llama_pos np = nPast;
        const int32_t rc = mtmd_helper_eval_chunk_single(mctx, ctx, audioChunk, nPast, /*seq_id*/ 0,
                                                        opt.nBatch, /*logits_last*/ false, &np);
        mtmd_bitmap_free(bmp);
        mtmd_input_chunks_free(chunks);
        if (rc != 0) { if (err) *err = "音频编码/解码失败"; return false; }

        nPast = np;
        return true;
    }
};

Qwen3Asr::Qwen3Asr() : d(new Impl) {}

Qwen3Asr::~Qwen3Asr()
{
    if (!d) return;
    if (d->mctx)  mtmd_free(d->mctx);
    if (d->ctx)   llama_free(d->ctx);
    if (d->model) llama_model_free(d->model);
    delete d;
}

int Qwen3Asr::sampleRate() const { return d->audioRate; }

bool Qwen3Asr::load(const AsrOptions & opt, std::string * error)
{
    const auto fail = [error](const std::string & msg) {
        if (error) *error = msg;
        return false;
    };

    d->opt = opt;

    if (!opt.verbose) {
        llama_log_set(silentLog, nullptr);
        mtmd_log_set(silentLog, nullptr);
    }

    llama_backend_init();

    llama_model_params mparams = llama_model_default_params();
    mparams.n_gpu_layers = 0;      // 本机无独显，强制 CPU
    mparams.use_mmap   = true;

    d->model = llama_model_load_from_file(opt.modelPath.c_str(), mparams);
    if (!d->model) {
        return fail("加载主模型失败: " + opt.modelPath);
    }
    d->vocab = llama_model_get_vocab(d->model);

    llama_context_params cparams = llama_context_default_params();
    cparams.n_ctx         = (uint32_t) opt.nCtx;
    cparams.n_batch       = (uint32_t) opt.nBatch;
    cparams.n_ubatch      = (uint32_t) opt.nUbatch;
    cparams.n_threads     = opt.threads;
    cparams.n_threads_batch = opt.threads;

    d->ctx = llama_init_from_model(d->model, cparams);
    if (!d->ctx) {
        return fail("创建 llama 上下文失败");
    }

    mtmd_context_params mctxParams = mtmd_context_params_default();
    mctxParams.use_gpu      = false;
    mctxParams.n_threads    = opt.threads;
    mctxParams.print_timings = false;
    mctxParams.warmup       = opt.warmup;

    d->mctx = mtmd_init_from_file(opt.mmprojPath.c_str(), d->model, mctxParams);
    if (!d->mctx) {
        return fail("加载 mmproj 失败: " + opt.mmprojPath);
    }

    if (!mtmd_support_audio(d->mctx)) {
        return fail("该 mmproj 不支持音频输入（需要 clip.projector_type = qwen3a 的版本）");
    }

    const int rate = mtmd_get_audio_sample_rate(d->mctx);
    if (rate > 0) d->audioRate = rate;

    return true;
}

Qwen3Asr::Result Qwen3Asr::transcribe(const float * pcm, size_t nSamples)
{
    return transcribeWithPrefix(pcm, nSamples, std::string());
}

std::string Qwen3Asr::dropLastTokens(const std::string & text, int dropLast) const
{
    if (!d->vocab || dropLast <= 0 || text.empty()) return text;

    std::vector<llama_token> tokens(text.size() + 16);
    int32_t n = llama_tokenize(d->vocab, text.c_str(), (int32_t) text.size(),
                               tokens.data(), (int32_t) tokens.size(),
                               /*add_special*/ false, /*parse_special*/ true);
    if (n < 0) {   // 缓冲区不够，按返回的所需长度重来
        tokens.resize((size_t) -n);
        n = llama_tokenize(d->vocab, text.c_str(), (int32_t) text.size(),
                           tokens.data(), (int32_t) tokens.size(), false, true);
        if (n < 0) return text;
    }
    tokens.resize((size_t) n);
    if ((int) tokens.size() <= dropLast) return {};

    tokens.resize(tokens.size() - (size_t) dropLast);

    // special=true：让 <asr_text> 这类特殊标记以字面形式保留，前缀才能原样回喂
    std::string out;
    char buf[512];
    for (llama_token id : tokens) {
        const int32_t m = llama_token_to_piece(d->vocab, id, buf, sizeof(buf), 0, /*special*/ true);
        if (m > 0) out.append(buf, (size_t) m);
    }
    return out;
}

Qwen3Asr::Result Qwen3Asr::transcribeWithPrefix(const float * pcm, size_t nSamples,
                                               const std::string & prefix)
{
    Result res;

    if (!d->model || !d->ctx || !d->mctx) {
        res.error = "模型未加载";
        return res;
    }
    if (nSamples == 0) {
        res.error = "音频为空";
        return res;
    }

    const auto t0 = std::chrono::steady_clock::now();

    // 每段独立转写，先清空 KV cache
    llama_memory_clear(llama_get_memory(d->ctx), true);

    const std::string marker = mtmd_default_marker();
    // 关键：assistant 头之后接上"已生成文本"，让模型从那里继续生成（前缀回滚）
    const std::string prompt = d->buildPrompt(marker) + prefix;

    mtmd_bitmap * bitmap = mtmd_bitmap_init_from_audio(nSamples, pcm);
    if (!bitmap) {
        res.error = "构造音频 bitmap 失败";
        return res;
    }

    mtmd_input_chunks * chunks = mtmd_input_chunks_init();
    const mtmd_input_text text{ prompt.c_str(), /*add_special*/ true, /*parse_special*/ true };
    const mtmd_bitmap * bitmaps[1] = { bitmap };

    const int32_t tokRes = mtmd_tokenize(d->mctx, chunks, &text, bitmaps, 1);
    if (tokRes != 0) {
        mtmd_bitmap_free(bitmap);
        mtmd_input_chunks_free(chunks);
        res.error = "mtmd_tokenize 失败 (code " + std::to_string(tokRes) + ")";
        return res;
    }

    llama_pos nPast = 0;
    const int32_t evalRes = mtmd_helper_eval_chunks(
        d->mctx, d->ctx, chunks,
        /*n_past*/ 0, /*seq_id*/ 0, /*n_batch*/ d->opt.nBatch,
        /*logits_last*/ true, &nPast);
    mtmd_bitmap_free(bitmap);
    if (evalRes != 0) {
        mtmd_input_chunks_free(chunks);
        res.error = "mtmd_helper_eval_chunks 失败 (code " + std::to_string(evalRes) + ")";
        return res;
    }
    mtmd_input_chunks_free(chunks);

    // 贪心解码
    llama_sampler * smpl = llama_sampler_chain_init(llama_sampler_chain_default_params());
    llama_sampler_chain_add(smpl, llama_sampler_init_greedy());

    std::string out;
    for (int i = 0; i < d->opt.maxTokens; ++i) {
        llama_token id = llama_sampler_sample(smpl, d->ctx, -1);
        if (llama_vocab_is_eog(d->vocab, id)) break;

        char buf[256];
        const int32_t n = llama_token_to_piece(d->vocab, id, buf, sizeof(buf), 0, /*special*/ true);
        if (n > 0) out.append(buf, (size_t) n);
        ++res.tokenCount;

        llama_batch batch = llama_batch_get_one(&id, 1);
        if (llama_decode(d->ctx, batch) != 0) {
            res.error = "llama_decode 失败";
            break;
        }
    }
    llama_sampler_free(smpl);

    const auto t1 = std::chrono::steady_clock::now();
    res.inferSeconds = std::chrono::duration<double>(t1 - t0).count();
    res.raw          = out;
    res.transcript   = parseAsrOutput(out);
    res.ok           = res.error.empty();
    return res;
}

// ---------------------------------------------------------------------------
// 流式（增量 KV 复用）
// ---------------------------------------------------------------------------

void Qwen3Asr::streamReset()
{
    if (!d->ctx) return;
    llama_memory_clear(llama_get_memory(d->ctx), true);
    d->streamReady   = false;
    d->headTokens.clear();
    d->tailTokens.clear();
    d->streamAudio.clear();
    d->stableSamples = 0;
    d->nPastStable   = 0;
}

bool Qwen3Asr::prepareStream(std::string * error)
{
    if (d->streamReady) return true;

    // 用一个占位音频 tokenize 一次，只为切出「音频之前」和「音频之后」的 prompt 文本 token。
    // 占位音频本身不会被解码。
    std::vector<float> dummy(1600, 0.0f);   // 0.1s
    mtmd_bitmap * bmp = mtmd_bitmap_init_from_audio(dummy.size(), dummy.data());
    mtmd_input_chunks * chunks = mtmd_input_chunks_init();

    const std::string marker = mtmd_default_marker();
    const std::string prompt = d->buildPrompt(marker);
    const mtmd_input_text itext{ prompt.c_str(), true, true };
    const mtmd_bitmap * bms[1] = { bmp };

    if (mtmd_tokenize(d->mctx, chunks, &itext, bms, 1) != 0) {
        mtmd_bitmap_free(bmp);
        mtmd_input_chunks_free(chunks);
        if (error) *error = "prepareStream: mtmd_tokenize 失败";
        return false;
    }
    mtmd_bitmap_free(bmp);

    bool seenAudio = false;
    for (size_t i = 0; i < mtmd_input_chunks_size(chunks); ++i) {
        const mtmd_input_chunk * c = mtmd_input_chunks_get(chunks, i);
        if (mtmd_input_chunk_get_type(c) == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
            seenAudio = true;
            continue;
        }
        size_t n = 0;
        const llama_token * toks = mtmd_input_chunk_get_tokens_text(c, &n);
        if (!toks || n == 0) continue;
        std::vector<llama_token> & dst = seenAudio ? d->tailTokens : d->headTokens;
        dst.insert(dst.end(), toks, toks + n);
    }
    mtmd_input_chunks_free(chunks);

    if (d->headTokens.empty() || d->tailTokens.empty()) {
        if (error) *error = "prepareStream: 没能切出音频前后的 prompt 文本";
        return false;
    }

    llama_memory_clear(llama_get_memory(d->ctx), true);
    if (!d->decodeTokens(d->headTokens.data(), (int32_t) d->headTokens.size(), 0, false)) {
        if (error) *error = "prepareStream: 解码 prompt 头部失败";
        return false;
    }
    d->nPastStable = (llama_pos) d->headTokens.size();
    d->streamReady = true;
    return true;
}

// 8 秒对齐的音频切片：这是编码器的注意力窗口（n_window_infer = 800 帧 @100Hz）。
// 按这个粒度切，编码结果才和「整段音频一次编码」一致；切小了会明显掉准确率。
static constexpr double kSliceSeconds = 8.0;

Qwen3Asr::StreamStep Qwen3Asr::streamAppend(const float * pcm, size_t nSamples,
                                            const std::string & prefix, int maxTokens)
{
    StreamStep res;
    if (!d->model || !d->ctx || !d->mctx) { res.error = "模型未加载"; return res; }

    std::string err;
    if (!prepareStream(&err)) { res.error = err; return res; }

    const auto t0 = std::chrono::steady_clock::now();

    d->streamAudio.insert(d->streamAudio.end(), pcm, pcm + nSamples);

    const size_t sliceSamples = (size_t) (kSliceSeconds * d->audioRate);
    const size_t targetStable = (d->streamAudio.size() / sliceSamples) * sliceSamples;

    const llama_pos nPastBefore = d->nPastStable;

    // 1) 已排满 8s 的切片固化进 KV，之后永远复用
    if (targetStable > d->stableSamples) {
        llama_memory_seq_rm(llama_get_memory(d->ctx), 0, d->nPastStable, -1);
        llama_pos np = d->nPastStable;
        if (!d->appendAudio(d->streamAudio.data() + d->stableSamples,
                            targetStable - d->stableSamples, np, &err)) {
            res.error = "固化音频切片失败: " + err;
            return res;
        }
        d->nPastStable   = np;
        d->stableSamples = targetStable;
    }

    // 2) 尾部不满 8s 的部分：每次重编码，但只是临时的（下一步会被丢掉）
    llama_memory_seq_rm(llama_get_memory(d->ctx), 0, d->nPastStable, -1);
    llama_pos nPast = d->nPastStable;
    if (d->streamAudio.size() > d->stableSamples) {
        if (!d->appendAudio(d->streamAudio.data() + d->stableSamples,
                            d->streamAudio.size() - d->stableSamples, nPast, &err)) {
            res.error = "编码尾部音频失败: " + err;
            return res;
        }
    }
    res.audioTokens = (int) (nPast - nPastBefore);
    const auto tEnc = std::chrono::steady_clock::now();

    // 3) 音频之后的 prompt 文本 + 前缀
    std::vector<llama_token> prefixToks;
    if (!prefix.empty()) {
        prefixToks.resize(prefix.size() + 16);
        int32_t n = llama_tokenize(d->vocab, prefix.c_str(), (int32_t) prefix.size(),
                                   prefixToks.data(), (int32_t) prefixToks.size(),
                                   /*add_special*/ false, /*parse_special*/ true);
        if (n < 0) {
            prefixToks.resize((size_t) -n);
            n = llama_tokenize(d->vocab, prefix.c_str(), (int32_t) prefix.size(),
                               prefixToks.data(), (int32_t) prefixToks.size(), false, true);
        }
        if (n < 0) { res.error = "前缀分词失败"; return res; }
        prefixToks.resize((size_t) n);
    }

    llama_pos p = nPast;
    const bool prefixLast = prefixToks.empty();
    if (!d->decodeTokens(d->tailTokens.data(), (int32_t) d->tailTokens.size(), p, prefixLast)) {
        res.error = "解码 tail 失败";
        return res;
    }
    p += (llama_pos) d->tailTokens.size();

    if (!prefixToks.empty()) {
        if (!d->decodeTokens(prefixToks.data(), (int32_t) prefixToks.size(), p, true)) {
            res.error = "解码前缀失败";
            return res;
        }
        p += (llama_pos) prefixToks.size();
    }

    // 4) 贪心生成（每个 token 的 logits 都要留给下一次采样）
    llama_sampler * smpl = llama_sampler_chain_init(llama_sampler_chain_default_params());
    llama_sampler_chain_add(smpl, llama_sampler_init_greedy());

    std::string out;
    for (int i = 0; i < maxTokens; ++i) {
        llama_token id = llama_sampler_sample(smpl, d->ctx, -1);
        if (llama_vocab_is_eog(d->vocab, id)) break;

        char buf[256];
        const int32_t n = llama_token_to_piece(d->vocab, id, buf, sizeof(buf), 0, /*special*/ true);
        if (n > 0) out.append(buf, (size_t) n);
        ++res.newTokens;

        if (!d->decodeTokens(&id, 1, p, true)) { res.error = "生成时解码失败"; break; }
        ++p;
    }
    llama_sampler_free(smpl);

    const auto t1 = std::chrono::steady_clock::now();
    res.encodeSeconds = std::chrono::duration<double>(tEnc - t0).count();
    res.inferSeconds  = std::chrono::duration<double>(t1 - t0).count();
    res.newText       = out;
    res.ok            = res.error.empty();
    return res;
}

} // namespace lc