#include "AapgsModelClassifier.h"

#include <cmath>
#include <cstdint>
#include <mutex>

#include <QCoreApplication>
#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QLibrary>
#include <QStringList>

#include "yolo_runtime.h"

namespace {

const char kModelConfigRelativePath[] =
    "models/primary/model.yaml";
const char kModelConfigFileName[] = "model.yaml";

void AppendUniquePath(QStringList& paths, const QString& path)
{
    const QString cleanPath = QDir::cleanPath(path.trimmed());
    if (!cleanPath.isEmpty() && !paths.contains(cleanPath)) {
        paths.push_back(cleanPath);
    }
}

QString EnvironmentPath(const char* name)
{
    return QString::fromLocal8Bit(qgetenv(name)).trimmed();
}

QString SourcePackageRoot()
{
    QDir sourceDir(QFileInfo(QString::fromUtf8(__FILE__)).absolutePath());
#ifdef Q_OS_WIN
    const QString packageDirectory = QStringLiteral("Windows");
#else
    const QString packageDirectory = QStringLiteral("Arm");
#endif
    return QDir::cleanPath(sourceDir.absoluteFilePath(
        QStringLiteral("../../../../../AppAlgo/AAPGS_model/%1")
            .arg(packageDirectory)));
}

QString PackageRootForConfig(const QString& configPath)
{
    QDir directory(QFileInfo(configPath).absolutePath());
    for (int level = 0; level < 5; ++level) {
        const bool hasRuntime =
            QFileInfo::exists(directory.filePath(
                QStringLiteral("lib/libyolo_runtime.so.5"))) ||
            QFileInfo::exists(directory.filePath(
                QStringLiteral("lib/yolo_runtime.dll")));
        if (hasRuntime) {
            return directory.absolutePath();
        }
        if (!directory.cdUp()) {
            break;
        }
    }
    return QFileInfo(configPath).absolutePath();
}

bool FindModelConfig(QString& configPath,
                     QString& packageRoot,
                     QString& errorMessage)
{
    const QString configuredPath = EnvironmentPath("AAPGS_MODEL_CONFIG");
    if (!configuredPath.isEmpty()) {
        QFileInfo configuredFile(configuredPath);
        if (configuredFile.isDir()) {
            const QDir configuredDir(configuredFile.absoluteFilePath());
            const QFileInfo directConfig(configuredDir.filePath(
                QString::fromLatin1(kModelConfigFileName)));
            configuredFile.setFile(directConfig.isFile()
                ? directConfig.absoluteFilePath()
                : configuredDir.filePath(QString::fromLatin1(
                      kModelConfigRelativePath)));
        }
        if (!configuredFile.isFile()) {
            errorMessage = QStringLiteral("AAPGS模型配置不存在：%1")
                               .arg(configuredPath);
            return false;
        }
        configPath = configuredFile.absoluteFilePath();
        packageRoot = PackageRootForConfig(configPath);
        return true;
    }

    QStringList roots;
    const QString configuredRoot = EnvironmentPath("AAPGS_MODEL_ROOT");
    if (!configuredRoot.isEmpty()) {
        AppendUniquePath(roots, configuredRoot);
    }

    const QString applicationDir = QCoreApplication::applicationDirPath();
    AppendUniquePath(roots, QDir(applicationDir).filePath(QStringLiteral("AAPGS_model")));
    AppendUniquePath(roots, QDir(applicationDir).filePath(QStringLiteral("aapgs")));
    AppendUniquePath(roots, QDir(applicationDir).filePath(QStringLiteral("../AAPGS_model")));
    AppendUniquePath(roots, QDir::current().filePath(QStringLiteral("AAPGS_model")));
    AppendUniquePath(roots, QDir::current().filePath(QStringLiteral("AppAlgo/AAPGS_model")));
    AppendUniquePath(roots, SourcePackageRoot());
    AppendUniquePath(roots, QStringLiteral("/opt/rk3588-ai/AAPGS_model"));
    AppendUniquePath(roots, QStringLiteral("/usr/lib/AAPGS_model"));
    AppendUniquePath(roots, QStringLiteral("/usr/local/lib/AAPGS_model"));

    for (const QString& root : roots) {
        const QFileInfo candidate(QDir(root).filePath(
            QString::fromLatin1(kModelConfigRelativePath)));
        if (candidate.isFile()) {
            configPath = candidate.absoluteFilePath();
            packageRoot = QDir(root).absolutePath();
            return true;
        }
    }

    errorMessage = QStringLiteral(
        "未找到AAPGS机型分类配置，请部署到"
        "/opt/rk3588-ai/AAPGS_model/"
        "models/primary/model.yaml，"
        "或设置AAPGS_MODEL_CONFIG");
    return false;
}

QString ModelName(const yolo_result_common_v1_t& item,
                  QString& errorMessage)
{
    int length = 0;
    while (length < YOLO_LABEL_LENGTH && item.label[length] != '\0') {
        ++length;
    }
    const QString label = QString::fromUtf8(item.label, length).trimmed();
    if (label.compare(QStringLiteral("a320"), Qt::CaseInsensitive) == 0 ||
        (label.isEmpty() && item.class_id == 0)) {
        return QStringLiteral("A320");
    }
    if (label.compare(QStringLiteral("b737"), Qt::CaseInsensitive) == 0 ||
        (label.isEmpty() && item.class_id == 1)) {
        return QStringLiteral("B737");
    }
    if (label.compare(QStringLiteral("other"), Qt::CaseInsensitive) == 0 ||
        label == QStringLiteral("其他") ||
        (label.isEmpty() && item.class_id == 2)) {
        return QStringLiteral("其他");
    }
    errorMessage = label.isEmpty()
        ? QStringLiteral("AAPGS返回未知机型类别：%1").arg(item.class_id)
        : QStringLiteral("AAPGS返回未知机型类别：%1/%2")
              .arg(label)
              .arg(item.class_id);
    return QString();
}

} // namespace

class AapgsModelClassifier::Impl
{
public:
    using GetAbiVersionFunction = uint32_t (*)();
    using CreateFunction = yolo_status_t (*)(const char*, yolo_runtime_t**,
                                              yolo_error_info_v1_t*);
    using GetInfoFunction = yolo_status_t (*)(const yolo_runtime_t*,
                                               yolo_runtime_info_v1_t*);
    using InferImageFunction = yolo_status_t (*)(
        yolo_runtime_t*, const yolo_image_view_v1_t*,
        const yolo_infer_options_v1_t*, yolo_result_set_t**,
        yolo_error_info_v1_t*);
    using ResultGetCountFunction = yolo_status_t (*)(
        const yolo_result_set_t*, uint32_t*);
    using ResultGetCommonFunction = yolo_status_t (*)(
        const yolo_result_set_t*, uint32_t, yolo_result_common_v1_t*);
    using ResultReleaseFunction = void (*)(yolo_result_set_t*);
    using DestroyFunction = void (*)(yolo_runtime_t*);

    ~Impl()
    {
        UnloadRuntimeLibrary();
    }

    bool Classify(const QImage& frame,
                  Classification& result,
                  QString& errorMessage)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        result = Classification();
        errorMessage.clear();

        if (frame.isNull() || frame.width() <= 0 || frame.height() <= 0) {
            errorMessage = QStringLiteral("AAPGS机型识别图像无效");
            return false;
        }
        if (!EnsureInitialized(errorMessage)) {
            return false;
        }

        const QImage rgbFrame = frame.format() == QImage::Format_RGB888
            ? frame
            : frame.convertToFormat(QImage::Format_RGB888);
        if (rgbFrame.isNull()) {
            errorMessage = QStringLiteral("AAPGS机型识别图像转换失败");
            return false;
        }

        yolo_image_view_v1_t image{};
        image.struct_size = sizeof(image);
        image.version = YOLO_IMAGE_VIEW_ABI_VERSION;
        image.data = rgbFrame.constBits();
        image.width = rgbFrame.width();
        image.height = rgbFrame.height();
        image.row_stride_bytes = rgbFrame.bytesPerLine();
        image.pixel_format = YOLO_PIXEL_FORMAT_RGB888;

        yolo_infer_options_v1_t options{};
        options.struct_size = sizeof(options);
        options.version = YOLO_INFER_OPTIONS_VERSION;
        options.seg_geometry_mode = YOLO_SEG_GEOMETRY_NONE;

        yolo_error_info_v1_t error = MakeErrorInfo();
        yolo_result_set_t* rawResults = nullptr;
        const yolo_status_t status = m_inferImage(
            m_runtime, &image, &options, &rawResults, &error);
        if (status != YOLO_STATUS_OK || !rawResults) {
            const QString runtimeError = RuntimeError(error);
            if (status == YOLO_STATUS_INITIALIZATION_FAILED ||
                status == YOLO_STATUS_BACKEND_FAILED) {
                UnloadRuntimeLibrary();
            }
            errorMessage = QStringLiteral("AAPGS模型推理失败(%1)：%2")
                               .arg(status)
                               .arg(runtimeError);
            return false;
        }

        const auto releaseResults = [this](yolo_result_set_t* results) {
            if (results && m_resultRelease) {
                m_resultRelease(results);
            }
        };
        std::unique_ptr<yolo_result_set_t, decltype(releaseResults)> results(
            rawResults, releaseResults);

        uint32_t resultCount = 0;
        if (m_resultGetCount(results.get(), &resultCount) != YOLO_STATUS_OK ||
            resultCount == 0 || resultCount > 4096) {
            errorMessage = QStringLiteral("AAPGS分类结果数量异常：%1")
                               .arg(resultCount);
            return false;
        }

        yolo_result_common_v1_t bestResult{};
        bool hasBestResult = false;
        for (uint32_t index = 0; index < resultCount; ++index) {
            yolo_result_common_v1_t item{};
            item.struct_size = sizeof(item);
            item.version = YOLO_RESULT_COMMON_VERSION;
            if (m_resultGetCommon(results.get(), index, &item) !=
                YOLO_STATUS_OK) {
                continue;
            }
            if (item.task != YOLO_TASK_CLS) {
                continue;
            }
            if (!std::isfinite(item.score) || item.score < 0.0f ||
                item.score > 1.0f) {
                continue;
            }
            if (!hasBestResult || item.score > bestResult.score) {
                bestResult = item;
                hasBestResult = true;
            }
        }
        if (!hasBestResult) {
            errorMessage = QStringLiteral("AAPGS未返回有效分类结果");
            return false;
        }
        result.modelType = ModelName(bestResult, errorMessage);
        result.confidence = static_cast<double>(bestResult.score);
        if (result.modelType.isEmpty()) {
            result = Classification();
            return false;
        }
        return true;
    }

private:
    bool EnsureInitialized(QString& errorMessage)
    {
        if (m_runtime) {
            return true;
        }

        QString configPath;
        QString packageRoot;
        if (!FindModelConfig(configPath, packageRoot, errorMessage) ||
            !LoadRuntimeLibrary(packageRoot, errorMessage)) {
            return false;
        }

        const uint32_t runtimeAbi = m_getAbiVersion();
        if (runtimeAbi != YOLO_RUNTIME_ABI_VERSION) {
            errorMessage = QStringLiteral("AAPGS运行库ABI不匹配：期望%1，实际%2")
                               .arg(YOLO_RUNTIME_ABI_VERSION)
                               .arg(runtimeAbi);
            UnloadRuntimeLibrary();
            return false;
        }

        const QByteArray encodedConfigPath = QFile::encodeName(configPath);
        yolo_error_info_v1_t error = MakeErrorInfo();
        const yolo_status_t status = m_create(
            encodedConfigPath.constData(), &m_runtime, &error);
        if (status != YOLO_STATUS_OK || !m_runtime) {
            const QString runtimeError = RuntimeError(error);
            UnloadRuntimeLibrary();
            errorMessage = QStringLiteral("AAPGS模型初始化失败(%1)：%2；配置：%3")
                               .arg(status)
                               .arg(runtimeError)
                               .arg(configPath);
            return false;
        }

        yolo_runtime_info_v1_t runtimeInfo{};
        runtimeInfo.struct_size = sizeof(runtimeInfo);
        runtimeInfo.version = YOLO_RUNTIME_INFO_VERSION;
        if (m_getInfo(m_runtime, &runtimeInfo) != YOLO_STATUS_OK) {
            UnloadRuntimeLibrary();
            errorMessage = QStringLiteral("读取AAPGS运行库信息失败");
            return false;
        }
        if (runtimeInfo.task != YOLO_TASK_CLS) {
            const int task = static_cast<int>(runtimeInfo.task);
            UnloadRuntimeLibrary();
            errorMessage = QStringLiteral("AAPGS模型任务不是分类任务：%1")
                               .arg(task);
            return false;
        }
        if ((runtimeInfo.capabilities & YOLO_CAPABILITY_INFER_IMAGE) == 0) {
            UnloadRuntimeLibrary();
            errorMessage = QStringLiteral("AAPGS运行库不支持内存图像推理");
            return false;
        }
        return true;
    }

    bool LoadRuntimeLibrary(const QString& packageRoot, QString& errorMessage)
    {
        if (m_library.isLoaded()) {
            if (ResolveFunctions()) {
                return true;
            }
            UnloadRuntimeLibrary();
        }

        QStringList candidates;
        const QString configuredLibrary = EnvironmentPath("AAPGS_RUNTIME_LIBRARY");
        if (!configuredLibrary.isEmpty()) {
            AppendUniquePath(candidates, configuredLibrary);
        }
#ifdef Q_OS_WIN
        AppendUniquePath(candidates, QDir(packageRoot).filePath(
            QStringLiteral("lib/yolo_runtime.dll")));
        AppendUniquePath(candidates, QStringLiteral("yolo_runtime"));
#else
        AppendUniquePath(candidates, QDir(packageRoot).filePath(
            QStringLiteral("lib/libyolo_runtime.so.5")));
        AppendUniquePath(candidates, QDir(packageRoot).filePath(
            QStringLiteral("lib/libyolo_runtime.so")));
        AppendUniquePath(candidates, QStringLiteral("libyolo_runtime.so.5"));
        AppendUniquePath(candidates, QStringLiteral("yolo_runtime"));
#endif

        QString lastLoadError;
        for (const QString& candidate : candidates) {
            m_library.setFileName(candidate);
            if (!m_library.load()) {
                lastLoadError = m_library.errorString();
                continue;
            }
            if (ResolveFunctions()) {
                return true;
            }
            lastLoadError = QStringLiteral("AAPGS运行库缺少必需的C ABI符号");
            UnloadRuntimeLibrary();
        }

        errorMessage = QStringLiteral("加载AAPGS运行库失败：%1")
                           .arg(lastLoadError);
        return false;
    }

    bool ResolveFunctions()
    {
        m_getAbiVersion = reinterpret_cast<GetAbiVersionFunction>(
            m_library.resolve("yolo_runtime_get_abi_version"));
        m_create = reinterpret_cast<CreateFunction>(
            m_library.resolve("yolo_runtime_create"));
        m_getInfo = reinterpret_cast<GetInfoFunction>(
            m_library.resolve("yolo_runtime_get_info"));
        m_inferImage = reinterpret_cast<InferImageFunction>(
            m_library.resolve("yolo_runtime_infer_image"));
        m_resultGetCount = reinterpret_cast<ResultGetCountFunction>(
            m_library.resolve("yolo_result_set_get_count"));
        m_resultGetCommon = reinterpret_cast<ResultGetCommonFunction>(
            m_library.resolve("yolo_result_set_get_common"));
        m_resultRelease = reinterpret_cast<ResultReleaseFunction>(
            m_library.resolve("yolo_result_set_release"));
        m_destroy = reinterpret_cast<DestroyFunction>(
            m_library.resolve("yolo_runtime_destroy"));
        return m_getAbiVersion && m_create && m_getInfo && m_inferImage &&
               m_resultGetCount && m_resultGetCommon && m_resultRelease &&
               m_destroy;
    }

    static yolo_error_info_v1_t MakeErrorInfo()
    {
        yolo_error_info_v1_t error{};
        error.struct_size = sizeof(error);
        error.version = YOLO_ERROR_INFO_VERSION;
        return error;
    }

    static QString RuntimeError(const yolo_error_info_v1_t& error)
    {
        int length = 0;
        while (length < static_cast<int>(sizeof(error.message)) &&
               error.message[length] != '\0') {
            ++length;
        }
        const QString text = QString::fromUtf8(error.message, length).trimmed();
        return text.isEmpty() ? QStringLiteral("未知错误") : text;
    }

    void UnloadRuntimeLibrary()
    {
        if (m_runtime && m_destroy) {
            m_destroy(m_runtime);
        }
        m_runtime = nullptr;
        if (m_library.isLoaded()) {
            m_library.unload();
        }
        m_getAbiVersion = nullptr;
        m_create = nullptr;
        m_getInfo = nullptr;
        m_inferImage = nullptr;
        m_resultGetCount = nullptr;
        m_resultGetCommon = nullptr;
        m_resultRelease = nullptr;
        m_destroy = nullptr;
    }

private:
    std::mutex m_mutex;
    QLibrary m_library;
    yolo_runtime_t* m_runtime = nullptr;
    GetAbiVersionFunction m_getAbiVersion = nullptr;
    CreateFunction m_create = nullptr;
    GetInfoFunction m_getInfo = nullptr;
    InferImageFunction m_inferImage = nullptr;
    ResultGetCountFunction m_resultGetCount = nullptr;
    ResultGetCommonFunction m_resultGetCommon = nullptr;
    ResultReleaseFunction m_resultRelease = nullptr;
    DestroyFunction m_destroy = nullptr;
};

AapgsModelClassifier::AapgsModelClassifier()
    : m_impl(std::make_unique<Impl>())
{
}

AapgsModelClassifier::~AapgsModelClassifier() = default;

bool AapgsModelClassifier::Classify(const QImage& frame,
                                    Classification& result,
                                    QString& errorMessage)
{
    return m_impl && m_impl->Classify(frame, result, errorMessage);
}