/*
 * Copyright 2024 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "speechrecognizer.h"

#include <jsoncpp/json/json.h>
#include <modelconfig/modelconfigmanager.h>

#include <fstream>
#include <utility>

#include "taskmanager/taskrunner.h"
#include "utils/logger.h"
#include "utils/util.h"

SpeechRecognizer::SpeechRecognizer() {}

SpeechRecognizer::~SpeechRecognizer() {
    destroy();
    LOGD("Speech recognizer destructor execution completed.");
}

void SpeechRecognizer::setRecognitionResultCallback(const RecognitionResultCallback& callback) {
    recognitionResultCallback_ = callback;
}

void SpeechRecognizer::setAudioConfig(const AudioConfigInfo& configInfo) {
    inputAudioInfo_ = configInfo.inputAudioInfo();
}

void SpeechRecognizer::startContinuousRecognition() {
    engineInUse_ = true;

    {
        std::lock_guard<std::mutex> lock(continuousModelStatusTrackerkMutex_);
        continuousModelStatusTracker_ = std::make_unique<ModelStatusTracker>(speechEngine_->currentModelName());
    }

    switch (inputAudioInfo_.inputSource) {
        case InputAudioInfo::InputSource::DEFAULT_MIC: {
            startContinuousRecognitionFromDefaultMic();
            break;
        }
        case InputAudioInfo::InputSource::STREAM: {
            startContinuousRecognitionFromAudioStream();
            break;
        }
        case InputAudioInfo::InputSource::PCM_FILE: {
            startContinuousRecognitionFromPcmFile(inputAudioInfo_.inputFilePath);
            break;
        }
        case InputAudioInfo::InputSource::RAW_DATA: {
            startContinuousRecognitionFromRawData(inputAudioInfo_.audioData);
            break;
        }
        default: {
            break;
        }
    }
}

void SpeechRecognizer::stopContinuousRecognition() {
    {
        std::make_unique<ModelStatusTracker>(speechEngine_->currentModelName());
        continuousModelStatusTracker_.reset();
    }
    if (inputAudioInfo_.inputSource == InputAudioInfo::InputSource::DEFAULT_MIC) {
        isRunningDefaultMic_ = false;
    }

    if (speechEngine_ != nullptr) {
        auto result = speechEngine_->stopContinuousRecognition();
        if (!result.success) {
            LOGE("Failed to call stopContinuousRecognition: {} {}.", result.error.getCode(), result.error.getMessage());
            callRecognitionResultCallback(makeResultFromEngineError(result.error));
        }
    }
    engineInUse_ = false;
}

bool SpeechRecognizer::init(int& errorCode, const ISpeechEnginePtr& speechEngine) {
    speechEngine_ = speechEngine;
    ModelStatusTracker speechModelStatusTracker(speechEngine_->currentModelName());
    auto resultContinuous = speechEngine_->initContinuousRecognitionModule();
    if (!resultContinuous.success) {
        LOGE("Failed to call initContinuousRecognitionModule: {} {}.", resultContinuous.error.getCode(),
             resultContinuous.error.getMessage());
    }

    auto resultOnce = speechEngine_->initRecognizeOnceModule();
    if (!resultOnce.success) {
        LOGE("Failed to call initRecognizeOnceModule: {} {}.", resultOnce.error.getCode(),
             resultOnce.error.getMessage());
    }

    if (!resultContinuous.success) {
        errorCode = engineError::getSpeechErrorCode(resultContinuous.error.getCode());
        callRecognitionResultCallback(makeResultFromEngineError(resultContinuous.error));
        return false;
    }

    if (!resultOnce.success) {
        errorCode = engineError::getSpeechErrorCode(resultOnce.error.getCode());
        callRecognitionResultCallback(makeResultFromEngineError(resultOnce.error));
        return false;
    }

    // 初始化回调函数
    std::shared_ptr sharedObject = getSharedObject();
    if (sharedObject == nullptr) {
        errorCode = AISDK_RUNTIME_ERROR;
        return false;
    }
    speechEngine_->setRecognitionResultCallback(
        [sharedObject](const ai_engine::core_ai::speech::RecognitionResult& result) {
            SpeechRecognitionResult recognitionResult{};
            recognitionResult.reason = int(result.reason);
            recognitionResult.text = result.text;
            recognitionResult.speakerId = result.speakerId;
            recognitionResult.errorCode = engineError::getSpeechErrorCode(result.error.getCode());
            recognitionResult.errorMessage = result.error.getMessage();
            sharedObject->callRecognitionResultCallback(recognitionResult);
            LOGD("Current model name：{} {}.", sharedObject->speechEngine_->engineName(), result.modelName);
            if (recognitionResult.errorCode != AISDK_NO_ERROR) {
                sharedObject->errorOccurred_ = true;
            }
        });
    return true;
}

void SpeechRecognizer::destroy() {
    stopContinuousRecognition();
    // 清理回调函数资源
    recognitionResultCallback_ = nullptr;
    if (speechEngine_ == nullptr) {
        return;
    }
    speechEngine_->setRecognitionResultCallback(nullptr);
    auto result = speechEngine_->destroyContinuousRecognitionModule();
    if (!result.success) {
        LOGE("Failed to destroy continuous recognition module: {} {}.", result.error.getCode(),
             result.error.getMessage());
    }
    result = speechEngine_->destroyRecognizeOnceModule();
    if (!result.success) {
        LOGE("Failed to destroy recognize once module: {} {}.", result.error.getCode(), result.error.getMessage());
    }
}

SpeechRecognitionResult SpeechRecognizer::makeResultFromEngineError(const ai_engine::EngineError& error) const {
    SpeechRecognitionResult result{};
    result.reason = int(ai_engine::core_ai::speech::SpeechResultReason::SpeechErrorOccurred);
    result.text = "";
    result.speakerId = -1;
    result.errorCode = engineError::getSpeechErrorCode(error.getCode());
    result.errorMessage = error.getMessage();
    return result;
}

void SpeechRecognizer::callRecognitionResultCallback(const SpeechRecognitionResult& result) const {
    if (recognitionResultCallback_ != nullptr) {
        recognitionResultCallback_(result);
    }

    if (result.errorCode > 0) {
        std::lock_guard<std::mutex> lock(continuousModelStatusTrackerkMutex_);
        if (continuousModelStatusTracker_) {
            continuousModelStatusTracker_.reset();
        }
    }
}

static std::vector<uint8_t> readAudioData(const std::string& filePath) {
    std::ifstream file(filePath, std::ios::binary);
    if (!file.is_open()) {
        LOGE("Failed to open file: {}.", filePath);
        return {};
    }

    file.seekg(0, std::ios::end);
    std::streampos fileSize = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<char> audioData(fileSize);
    file.read(audioData.data(), fileSize);
    return {audioData.begin(), audioData.end()};
}

void SpeechRecognizer::startContinuousRecognitionFromDefaultMic() {
    errorOccurred_ = false;
    if (!AudioRecorder::isMicrophoneAvailable()) {
        SpeechRecognitionResult result;
        result.reason = int(ai_engine::core_ai::speech::SpeechResultReason::SpeechErrorOccurred);
        result.text = "";
        result.speakerId = -1;
        result.errorCode = SPEECH_DEFAULT_MICROPHONE_INVALID;
        result.errorMessage = "Default microphone invalid.";
        callRecognitionResultCallback(result);
        return;
    }

    isRunningDefaultMic_ = true;
    std::shared_ptr audioRecorderPtr{std::make_shared<AudioRecorder>()};
    std::shared_ptr sharedObject = getSharedObject();

    task::TaskRunner::getInstance().addTask({task::generateTaskId(), task::TaskType::ContinuousSpeechRecognition,
                                             [sharedObject, audioRecorderPtr{std::move(audioRecorderPtr)}]() {
                                                 sharedObject->runTaskForContinuousRecognitionFromDefaultMic(
                                                     audioRecorderPtr);
                                             }});
}

void SpeechRecognizer::runTaskForContinuousRecognitionFromDefaultMic(
    const std::shared_ptr<AudioRecorder>& audioRecorderPtr) const {
    auto result = speechEngine_->startContinuousRecognition(R"({"format": "", "rate": 16000, "channel": 1})");
    if (!result.success) {
        LOGE("Failed to call startContinuousRecognition: {} {}.", result.error.getCode(), result.error.getMessage());
        callRecognitionResultCallback(makeResultFromEngineError(result.error));
        return;
    }
    audioRecorderPtr->start();
    while (isRunningDefaultMic_) {
        using namespace std::chrono_literals;
        std::chrono::time_point<std::chrono::steady_clock> timePoint = std::chrono::steady_clock::now() + 40ms;
        if (errorOccurred_) {
            // 回调中报错，停止识别
            break;
        }
        std::string micName = AudioRecorder::getMicrophoneName();
        if (micName.empty()) {
            SpeechRecognitionResult result{};
            result.reason = int(ai_engine::core_ai::speech::SpeechResultReason::SpeechErrorOccurred);
            result.text = "";
            result.speakerId = -1;
            result.errorCode = SPEECH_DEFAULT_MICROPHONE_INVALID;
            result.errorMessage = "Default microphone invalid.";
            callRecognitionResultCallback(result);
            break;
        }

        std::vector<uint8_t> audioData = audioRecorderPtr->getSingleFrameAudioData();
        if (audioData.empty()) {
            std::this_thread::sleep_for(10ms);
            continue;
        }
        result = speechEngine_->writeContinuousRecognitionAudioData(audioData);
        if (!result.success) {
            LOGE("Failed to write audio data: {} {}.", result.error.getCode(), result.error.getMessage());
            callRecognitionResultCallback(makeResultFromEngineError(result.error));
            break;
        }
        std::this_thread::sleep_until(timePoint);
    }
}

void SpeechRecognizer::startContinuousRecognitionFromRawData(const std::vector<uint8_t>& audioData) {
    if (audioData.size() < SINGLE_FRAME_LIMIT) {
        LOGE("Func startContinuousRecognitionFromRawData params is abnormal!");
        ai_engine::EngineError error(int(ai_engine::SpeechEngineErrorCode::TooLittleData),
                                     "Input raw data is too little!");
        callRecognitionResultCallback(makeResultFromEngineError(error));
        return;
    }
    std::shared_ptr sharedObject = getSharedObject();

    task::TaskRunner::getInstance().addTask(
        {task::generateTaskId(), task::TaskType::ContinuousSpeechRecognition,
         [sharedObject, audioData]() { sharedObject->runTaskForContinuousRecognitionFromRawData(audioData); }});
}

void SpeechRecognizer::runTaskForContinuousRecognitionFromRawData(const std::vector<uint8_t>& audioData) {
    auto result = speechEngine_->startContinuousRecognition(R"({"format": "", "rate": 16000, "channel": 1})");

    if (!result.success) {
        LOGE("Failed to call startContinuousRecognition: {} {}.", result.error.getCode(), result.error.getMessage());
        callRecognitionResultCallback(makeResultFromEngineError(result.error));
    }
    for (int i = 0; i < audioData.size(); i += 1280) {
        LOGD("write audio data: {}.", i);
        auto begin = audioData.begin() + i;
        auto end = audioData.begin() + i + 1280;
        if (end > audioData.end()) {
            end = audioData.end();
        }
        result = speechEngine_->writeContinuousRecognitionAudioData(std::vector(begin, end));
        if (!result.success) {
            LOGE("Failed to write audio data: {} {}.", result.error.getCode(), result.error.getMessage());
            callRecognitionResultCallback(makeResultFromEngineError(result.error));
            break;
        }
    }
    result = speechEngine_->stopContinuousRecognition();
    if (!result.success) {
        LOGE("Failed to call stopContinuousRecognition: {} {}.", result.error.getCode(), result.error.getMessage());
        callRecognitionResultCallback(makeResultFromEngineError(result.error));
    }
}

void SpeechRecognizer::startContinuousRecognitionFromAudioStream() {
    auto result = speechEngine_->startContinuousRecognition(R"({"format": "", "rate": 16000, "channel": 1})");

    if (!result.success) {
        LOGE("Failed to call startContinuousRecognition: {} {}.", result.error.getCode(), result.error.getMessage());
        callRecognitionResultCallback(makeResultFromEngineError(result.error));
    }
}

void SpeechRecognizer::writeContinuousRecognitionFromAudioStream(const std::vector<uint8_t>& audioData) {
    std::shared_ptr sharedObject = getSharedObject();

    task::TaskRunner::getInstance().addTask(
        {task::generateTaskId(), task::TaskType::ContinuousSpeechRecognition, [sharedObject, audioData]() {
             auto result = sharedObject->speechEngine_->writeContinuousRecognitionAudioData(audioData);

             if (!result.success) {
                 LOGE("Failed to write audio data: {} {}.", result.error.getCode(), result.error.getMessage());
                 sharedObject->callRecognitionResultCallback(sharedObject->makeResultFromEngineError(result.error));
             }
         }});
}

bool SpeechRecognizer::isRegularPcmFile(const std::filesystem::path& path) {
    if (not std::filesystem::exists(path)) {
        LOGE("{} is not exists.", path.string());
        ai_engine::EngineError error(int(ai_engine::SpeechEngineErrorCode::InvalidParamType),
                                     fmt::format("Input path {} is not exists!", path.string()));
        callRecognitionResultCallback(makeResultFromEngineError(error));
        return false;
    }
    if (not std::filesystem::is_regular_file(path)) {
        LOGE("{} is not regular file.", path.string());
        ai_engine::EngineError error(int(ai_engine::SpeechEngineErrorCode::InvalidParamType),
                                     fmt::format("Input path {} is not exists!", path.string()));
        callRecognitionResultCallback(makeResultFromEngineError(error));
    }

    if (std::filesystem::file_size(path) < SINGLE_FRAME_LIMIT) {
        LOGE("file[{}] size < {}, it is too small.", path.string(), SINGLE_FRAME_LIMIT);
        ai_engine::EngineError error(int(ai_engine::SpeechEngineErrorCode::TooLittleData),
                                     fmt::format("Input path {} data is too little!", path.string()));
        callRecognitionResultCallback(makeResultFromEngineError(error));
        return false;
    }
    return true;
}

void SpeechRecognizer::startContinuousRecognitionFromPcmFile(const std::filesystem::path& path) {
    if (not isRegularPcmFile(path)) {
        return;
    }
    // TODO 这里需要优化一下，不会一次性读取
    startContinuousRecognitionFromRawData(readAudioData(path.string()));
}

void SpeechRecognizer::recognizeOnce() {
    audioInfoQueue_.enqueue(inputAudioInfo_);
    onceAsrAsync();
}

void SpeechRecognizer::onceAsrAsync() {
    if (taskId_ != task::INVALID_TASK_ID) {
        return;
    }
    taskId_ = task::generateTaskId();
    std::shared_ptr sharedObject = getSharedObject();
    task::TaskRunner::getInstance().addTask(
        {task::generateTaskId(), task::TaskType::OnceSpeechRecognition, [sharedObject]() {
             sharedObject->onceAsrTask();
             sharedObject->taskId_ = task::INVALID_TASK_ID;
         }});
}

void SpeechRecognizer::onceAsrTask() {
    InputAudioInfo inputAudioInfo;
    while (audioInfoQueue_.wait_dequeue_timed(inputAudioInfo, 10 * 1000 * 1000)) {
        Json::Value params;
        params["format"] = inputAudioInfo.getFormaString();
        params["rate"] = inputAudioInfo.rate;
        params["channel"] = inputAudioInfo.channel;
        std::vector<uint8_t> audioData;
        if (inputAudioInfo.inputSource == InputAudioInfo::InputSource::RAW_DATA) {
            audioData = inputAudioInfo.audioData;
        } else if (inputAudioInfo.inputSource == InputAudioInfo::InputSource::PCM_FILE) {
            audioData = readAudioData(inputAudioInfo.inputFilePath.string());
        } else {
            ai_engine::EngineError error(int(ai_engine::SpeechEngineErrorCode::InvalidParamType),
                                         "Recognize once only support file and raw data.");
            callRecognitionResultCallback(makeResultFromEngineError(error));
            return;
        }
        ModelStatusTracker modelStatusTracker(speechEngine_->currentModelName());
        engineInUse_ = true;
        auto result = speechEngine_->recognizeOnce(util::jsonToString(params), audioData);
        if (!result.success) {
            LOGE("Failed to speech recognize once: {} {}.", result.error.getCode(), result.error.getMessage());
            callRecognitionResultCallback(makeResultFromEngineError(result.error));
        }
        engineInUse_ = false;
    }
}
