#include "TTSManager.h"
#include <QDir>
#include <QFile>
#include <QDebug>

TTSManager::TTSManager(QObject* parent) : QObject(parent) {
}

TTSManager::~TTSManager() {
    cleanup();
}

bool TTSManager::initialize() {
    QString dataPath = QDir::homePath() + "/.config/QSmartAssistant/Data/";
    
    // 初始化TTS模型 - 支持中英文混合
    // 优先尝试使用新的vits-melo-tts-zh_en模型
    QString ttsModelPath = dataPath + "vits-melo-tts-zh_en/model.int8.onnx";
    QString ttsLexiconPath = dataPath + "vits-melo-tts-zh_en/lexicon.txt";
    QString ttsTokensPath = dataPath + "vits-melo-tts-zh_en/tokens.txt";
    QString ttsDataDirPath = "";  // melo-tts模型不需要额外的数据目录
    QString ttsDictDirPath = dataPath + "vits-melo-tts-zh_en/dict";  // jieba字典目录

    // 如果新模型不存在，尝试其他中英文模型
    if (!QFile::exists(ttsModelPath)) {
        ttsModelPath = dataPath + "vits-zh-en/vits-zh-en.onnx";
        ttsLexiconPath = dataPath + "vits-zh-en/lexicon.txt";
        ttsTokensPath = dataPath + "vits-zh-en/tokens.txt";
        ttsDataDirPath = dataPath + "vits-zh-en/espeak-ng-data";
        ttsDictDirPath = "";
    }

    // 最后回退到中文模型
    if (!QFile::exists(ttsModelPath)) {
        ttsModelPath = dataPath + "vits-zh-aishell3/vits-aishell3.int8.onnx";
        ttsLexiconPath = dataPath + "vits-zh-aishell3/lexicon.txt";
        ttsTokensPath = dataPath + "vits-zh-aishell3/tokens.txt";
        ttsDataDirPath = "";
        ttsDictDirPath = "";
    }

    currentModelPath = ttsModelPath;

    memset(&ttsConfig, 0, sizeof(ttsConfig));
    ttsConfig.model.vits.noise_scale = 0.667;
    ttsConfig.model.vits.noise_scale_w = 0.8;
    ttsConfig.model.vits.length_scale = 1.0;
    ttsConfig.model.num_threads = 2;
    ttsConfig.model.provider = "cpu";

    ttsModelPathStd = ttsModelPath.toStdString();
    ttsLexiconPathStd = ttsLexiconPath.toStdString();
    ttsTokensPathStd = ttsTokensPath.toStdString();
    ttsDataDirPathStd = ttsDataDirPath.toStdString();
    ttsDictDirPathStd = ttsDictDirPath.toStdString();

    ttsConfig.model.vits.model = ttsModelPathStd.c_str();
    ttsConfig.model.vits.lexicon = ttsLexiconPathStd.c_str();
    ttsConfig.model.vits.tokens = ttsTokensPathStd.c_str();

    // 如果有espeak数据目录，设置它以支持英文发音
    if (!ttsDataDirPath.isEmpty()) {
        ttsConfig.model.vits.data_dir = ttsDataDirPathStd.c_str();
    }

    // 如果有jieba字典目录，设置它以支持中文分词
    if (!ttsDictDirPath.isEmpty()) {
        ttsConfig.model.vits.dict_dir = ttsDictDirPathStd.c_str();
    }

    ttsSynthesizer = const_cast<SherpaOnnxOfflineTts*>(SherpaOnnxCreateOfflineTts(&ttsConfig));

    qDebug() << "TTS合成器:" << (ttsSynthesizer ? "成功" : "失败");
    qDebug() << "TTS模型类型:" << getModelType();

    return ttsSynthesizer != nullptr;
}

bool TTSManager::synthesizeText(const QString& text, int speakerId, const QString& outputPath) {
    if (!ttsSynthesizer) {
        return false;
    }

    // 生成音频
    const SherpaOnnxGeneratedAudio* audio = SherpaOnnxOfflineTtsGenerate(
        ttsSynthesizer, text.toUtf8().constData(), speakerId, 1.0);

    if (!audio || audio->n == 0) {
        if (audio) SherpaOnnxDestroyOfflineTtsGeneratedAudio(audio);
        return false;
    }

    // 保存为WAV文件
    bool success = SherpaOnnxWriteWave(audio->samples, audio->n, audio->sample_rate,
                                      outputPath.toUtf8().constData());

    SherpaOnnxDestroyOfflineTtsGeneratedAudio(audio);
    return success;
}

QString TTSManager::getModelType() const {
    if (currentModelPath.contains("vits-melo-tts-zh_en")) {
        return "MeloTTS中英文混合模型";
    } else if (currentModelPath.contains("vits-zh-en")) {
        return "VITS中英文混合模型";
    } else {
        return "中文模型";
    }
}

bool TTSManager::isMultilingualModel() const {
    return currentModelPath.contains("vits-melo-tts-zh_en") || 
           currentModelPath.contains("vits-zh-en");
}

void TTSManager::cleanup() {
    // 清理TTS合成器
    if (ttsSynthesizer) {
        SherpaOnnxDestroyOfflineTts(ttsSynthesizer);
        ttsSynthesizer = nullptr;
    }
}