/*
 * 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 "textrecognition.h"

#include <jsoncpp/json/json.h>

#include <filesystem>
#include <fstream>
#include <utility>

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

TextRecognition::TextRecognition(
    std::unique_ptr<ai_engine::core_ai::vision::AbstractTextRecognitionEngine> visionEngine)
    : visionEngine_(std::move(visionEngine)) {}

TextRecognition::~TextRecognition() { LOGD("~TextRecognition"); }

std::vector<TextLine> convertResult2Textlines(const ai_engine::core_ai::vision::TextRecognitionResult& result) {
    std::vector<TextLine> textLines;
    textLines.reserve(result.lines.size());
    for (const auto& line : result.lines) {
        std::vector<PixelPoint> points;
        points.reserve(line.points.size());
        for (const auto& point : line.points) {
            points.push_back({point.x, point.y});
        }
        textLines.push_back({line.text, points, line.conf});
    }
    return textLines;
}

void TextRecognition::recognizeTextFromImageFile(const std::string& imageFilePath,
                                                 RecognitionResultCallback resultCallback) {
    if (resultCallback == nullptr) {
        return;
    }
    const auto& request = generateTextRecognitionRequestWithFilePath(imageFilePath, std::move(resultCallback));
    textRecognitionRequests_.enqueue(request);
    textRecogniseAsync();
}

void TextRecognition::recognizeTextFromImageData(const std::string& imageData,
                                                 RecognitionResultCallback resultCallback) {
    if (resultCallback == nullptr) {
        return;
    }
    const auto& request = generateTextRecognitionRequestWithBase64Data(imageData, std::move(resultCallback));
    textRecognitionRequests_.enqueue(request);
    textRecogniseAsync();
}

TextRecognition::TextRecognitionRequest TextRecognition::generateTextRecognitionRequestWithFilePath(
    const std::string& imageFilePath, RecognitionResultCallback resultCallback) {
    TextRecognitionRequest request{};
    request.resultCallback = std::move(resultCallback);
    request.type = TextRecognition::TextRecognitionRequest::Type::ImagePath;
    request.imagePath = imageFilePath;
    return request;
}

TextRecognition::TextRecognitionRequest TextRecognition::generateTextRecognitionRequestWithBase64Data(
    const std::string& base64Data, RecognitionResultCallback resultCallback) {
    TextRecognitionRequest request{};
    request.resultCallback = std::move(resultCallback);
    request.type = TextRecognition::TextRecognitionRequest::Type::Base64Data;
    request.base64Data = base64Data;
    return request;
}

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

void TextRecognition::textRecognizeTask() {
    while (textRecognitionRequests_.wait_dequeue_timed(currentRequest_, 10 * 1000 * 1000)) {
        switch (currentRequest_.type) {
            case TextRecognitionRequest::Type::Base64Data: {
                visionEngine_->recognizeTextFromBase64ImageData(currentRequest_.base64Data);
                break;
            }
            case TextRecognitionRequest::Type::ImagePath: {
                visionEngine_->recognizeTextFromFile(currentRequest_.imagePath);
                break;
            }
            default: {
                LOGE("Unknown input type.");
                break;
            }
        }
    }
}

bool TextRecognition::init() {
    auto result = visionEngine_->init();

    if (!result.success) {
        LOGE("Failed to init vision engine: {} {}", result.error.getCode(), result.error.getMessage());
        callTextRecognitionResultCallback(makeResultFromEngineError(result.error));
        return false;
    }

    std::shared_ptr sharedObject = getSharedObject();
    if (sharedObject == nullptr) {
        return false;
    }
    // 该回调函数中捕获了shared_from_this的智能指针，因此要主动清空回调函数（destroy函数）
    visionEngine_->setTextRecognitionResultCallback(
        [sharedObject](const ai_engine::core_ai::vision::TextRecognitionResult& result) {
            TextRecognitionResult recognitionResult{result.allText, convertResult2Textlines(result),
                                                    visionEngineError::getCoreVisionErrorCode(result.error.getCode()),
                                                    result.error.getMessage()};
            sharedObject->callTextRecognitionResultCallback(recognitionResult);
        });

    return true;
}

void TextRecognition::destroy() const {
    if (visionEngine_ != nullptr) {
        visionEngine_->setTextRecognitionResultCallback(nullptr);
    }
}

void TextRecognition::callTextRecognitionResultCallback(const TextRecognitionResult& result) const {
    if (currentRequest_.resultCallback != nullptr) {
        currentRequest_.resultCallback(result);
    }
}

TextRecognitionResult TextRecognition::makeResultFromEngineError(const ai_engine::EngineError& error) const {
    return TextRecognitionResult{
        "", {}, visionEngineError::getCoreVisionErrorCode(error.getCode()), error.getMessage()};
}