#include "arg.h"
#include "common.h"
#include "sampling.h"
#include "log.h"
#include "llama.h"
#include "mtmd.h"
#include "mtmd-helper.h"

#include <cstdio>
#include <cstring>
#include <string>

/**
 * Please note that this is NOT a production-ready binary.
 * It is a playground for trying TTS support in llama.cpp.
 * For contributors: please keep this code simple and easy to understand. Do not add unnecessary complexity. The goal is to have a simple CLI for testing TTS support.
 */

static void print_usage(int, char ** argv) {
    LOG("\nexample usage:\n");
    LOG("\n    %s -m backbone.gguf -mm mmproj.gguf -p \"text to speak\" -o output.wav", argv[0]);
    LOG("\n    %s -hf user/model -p \"text to speak\" -o output.wav\n", argv[0]);
    LOG("\nnote: --tts-lang and --tts-speaker-file may not be supported in all models");
    LOG("\n      use -n to limit the output length");
    LOG("\n\n");
}

int main(int argc, char ** argv) {
    common_params params;

    common_init();

    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_TTS, print_usage)) {
        return 1;
    }

    mtmd_helper_log_set(common_log_default_callback, nullptr);

    if (params.prompt.empty()) {
        LOG_ERR("no prompt provided, use -p \"text\"\n");
        return 1;
    }
    if (params.mmproj.path.empty()) {
        LOG_ERR("no mmproj provided, use --mmproj\n");
        return 1;
    }

    llama_backend_init();
    llama_numa_init(params.numa);

    auto llama_init = common_init_from_params(params);
    llama_model    * model = llama_init->model();
    llama_context  * lctx  = llama_init->context();
    common_sampler * smpl  = llama_init->sampler(0);
    if (!model || !lctx) {
        LOG_ERR("failed to init model/context\n");
        return 1;
    }

    mtmd_context_params mtmd_params = mtmd_context_params_default();
    mtmd_params.use_gpu = params.mmproj_use_gpu;
    mtmd::context_ptr mctx(mtmd_init_from_file(params.mmproj.path.c_str(), model, mtmd_params));
    if (!mctx) {
        LOG_ERR("failed to load mmproj %s\n", params.mmproj.path.c_str());
        return 1;
    }
    if (mtmd_gen_audio_get_info(mctx.get()).type == MTMD_GEN_AUDIO_TYPE_NONE) {
        LOG_ERR("mmproj does not support audio generation\n");
        return 1;
    }

    mtmd::bitmap_ptr speaker_bitmap;
    if (!params.tts_speaker_file.empty()) {
        auto wrapper = mtmd_helper_bitmap_init_from_file(mctx.get(), params.tts_speaker_file.c_str(), false);
        if (!wrapper.bitmap) {
            LOG_ERR("failed to load speaker file %s\n", params.tts_speaker_file.c_str());
            return 1;
        }
        speaker_bitmap.reset(wrapper.bitmap);
    }

    mtmd_helper::gen_audio gen(lctx, mctx.get());
    mtmd_helper_gen_audio_inp inp{};
    inp.prompt      = params.prompt.c_str();
    inp.prompt_len  = params.prompt.size();
    inp.speaker_ref = speaker_bitmap.get();
    inp.lang        = params.tts_lang.c_str();
    inp.top_k       = params.sampling.top_k;
    inp.top_p       = params.sampling.top_p;
    inp.out_type    = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
    if (gen.set_input(&inp) != 0) {
        LOG_ERR("set_input failed\n");
        return 1;
    }

    // codec_0 (backbone) EOS token: ordinary LLM sampling concern, kept out of the
    // model-agnostic audio-generation helper
    const llama_vocab * vocab = llama_model_get_vocab(model);
    llama_token codec_eos_tok = LLAMA_TOKEN_NULL;
    for (llama_token t = 0; t < llama_vocab_n_tokens(vocab); t++) {
        if (!strcmp(llama_vocab_get_text(vocab, t), "<|codec_eos_token|>")) { codec_eos_tok = t; break; }
    }
    if (codec_eos_tok == LLAMA_TOKEN_NULL) {
        LOG_ERR("missing codec eos token in vocab\n");
        return 1;
    }

    auto sample_codec0 = [&]() -> llama_token {
        llama_token t = common_sampler_sample(smpl, lctx, -1);
        common_sampler_accept(smpl, t, true);
        return t;
    };

    const int max_new = params.n_predict > 0 ? params.n_predict : 512;
    int n_frames = 0;
    llama_token sampled = sample_codec0();
    const float * h_state = llama_get_embeddings_ith(lctx, -1);
    for (; n_frames < max_new && sampled != codec_eos_tok; n_frames++) {
        const float * h_next = nullptr;
        if (gen.step(sampled, h_state, &h_next) != 0) {
            LOG_ERR("step failed at frame %d\n", n_frames);
            return 1;
        }
        h_state = h_next;
        sampled = sample_codec0();
    }

    int32_t      sample_rate = 0;
    const char * data        = nullptr;
    size_t       data_len    = 0;
    if (gen.get_output(&sample_rate, &data, &data_len) != 0) {
        LOG_ERR("get_output failed\n");
        return 1;
    }

    LOG_INF("generated %d frames, %zu bytes of WAV audio (%d Hz)\n", n_frames, data_len, sample_rate);
    FILE * f = fopen(params.out_file.c_str(), "wb");
    if (!f) {
        LOG_ERR("failed to open %s\n", params.out_file.c_str());
        return 1;
    }
    fwrite(data, 1, data_len, f);
    fclose(f);
    LOG_INF("wrote %s\n", params.out_file.c_str());

    llama_backend_free();
    return 0;
}