#include "inflect_tts.hpp"
#include "ax_engine.h"
#include "ax_sys.h"
#include <cstring>
#include <stdexcept>
#include <cmath>
#include <algorithm>

// ---------------------------------------------------------------------------
// TTSEngine — NPU lifecycle
// ---------------------------------------------------------------------------

TTSEngine::TTSEngine(const std::string& encoder_path,
                     const std::string& decoder_path) {
    auto* eng = ax_engine_init();
    if (!eng)
        throw std::runtime_error("ax_engine_init failed");
    enc_engine_ = ax_engine_load_model(eng, encoder_path.c_str());
    dec_engine_ = ax_engine_load_model(eng, decoder_path.c_str());
    if (!enc_engine_ || !dec_engine_)
        throw std::runtime_error("Failed to load AX models");
}

TTSEngine::~TTSEngine() {
    if (enc_context_) ax_engine_destroy_context(enc_context_);
    if (dec_context_) ax_engine_destroy_context(dec_context_);
}

// ---------------------------------------------------------------------------
// encode — run encoder NPU inference
//   tokens [token_len]  →  4 output tensors {m_p, logs_p, x, x_mask}
// ---------------------------------------------------------------------------
std::vector<float> TTSEngine::encode(const int64_t* tokens, int token_len) {
    // 1) Embedding placeholder — in production, port PyTorch emb via libtorch.
    //    For now, zeros (caller must supply pre-computed x_emb).
    (void)tokens; (void)token_len;

    std::vector<float> empty;
    return empty;  // caller fills via direct x_emb
}

// ---------------------------------------------------------------------------
// decode — run decoder NPU inference
//   z_p [kInterChannels * mel_len]  →  waveform [mel_len * kHopLength]
// ---------------------------------------------------------------------------
std::vector<float> TTSEngine::decode(const float* z_p, int mel_len) {
    auto* engine = static_cast<ax_engine_t*>(dec_engine_);
    ax_engine_io_t io;
    ax_engine_get_io(engine, &io);

    // Zero-pad inputs to fixed shapes
    std::vector<float> z_p_padded(kInterChannels * kMaxMelFrames, 0.0f);
    std::vector<float> y_mask_padded(kMaxMelFrames, 0.0f);
    for (int i = 0; i < std::min(mel_len, kMaxMelFrames); ++i) {
        for (int c = 0; c < kInterChannels; ++c)
            z_p_padded[c * kMaxMelFrames + i] = z_p[c * mel_len + i];
        y_mask_padded[i] = 1.0f;
    }

    std::memcpy(io.inputs[0].data, z_p_padded.data(),
                kInterChannels * kMaxMelFrames * sizeof(float));
    std::memcpy(io.inputs[1].data, y_mask_padded.data(),
                kMaxMelFrames * sizeof(float));

    ax_engine_run(engine, &io);

    int out_len = std::min(mel_len, kMaxMelFrames) * kHopLength;
    std::vector<float> waveform(out_len);
    std::memcpy(waveform.data(), io.outputs[0].data, out_len * sizeof(float));
    return waveform;
}

// ---------------------------------------------------------------------------
// synthesize — full pipeline
// ---------------------------------------------------------------------------
std::vector<float> TTSEngine::synthesize(const std::string& text,
                                          float speed, float variation) {
    // Full pipeline requires:
    //   1. Text → phonemes → token IDs          (CPU, Python frontend port)
    //   2. Token IDs → embedding                (CPU, libtorch)
    //   3. Embedding → NPU encoder → stats      (NPU)
    //   4. Duration prediction + alignment      (CPU, libtorch)
    //   5. z_p → NPU decoder → waveform         (NPU)
    //
    // For a complete C++ pipeline, build with libtorch and the
    // Python frontend logic ported to C++. See model_convert/ for
    // ONNX sources and compile configs.
    //
    // This skeleton demonstrates the NPU API — plug in your text
    // processor and return raw waveform.

    (void)text; (void)speed; (void)variation;
    return {};  // skeleton — fill with your full pipeline
}