#include "AapgsModelClassifier.h"

#include <cmath>
#include <mutex>

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

#include "yolo_runtime.h"

namespace {

const char kModelConfigRelativePath[] =
    "models/model_20260712T020258_0000_6339d04b/"
    "model_20260712T020258_0000_6339d04b.yaml";
const char kModelConfigFileName[] =
    "model_20260712T020258_0000_6339d04b.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());
    return QDir::cleanPath(sourceDir.absoluteFilePath(
        QStringLiteral("../../../../../AppAlgo/AAPGS_model")));
}

QString PackageRootForConfig(const QString& configPath)
{
    QDir directory(QFileInfo(configPath).absolutePath());
    for (int level = 0; level < 5; ++level) {
        if (QFileInfo::exists(directory.filePath(
                QStringLiteral("lib/libyolo_runtime.so.3")))) {
            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/model_20260712T020258_0000_6339d04b/"
        "model_20260712T020258_0000_6339d04b.yaml，"
        "或设置AAPGS_MODEL_CONFIG");
    return false;
}

QString ModelName(const yolo_result_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");
    }
    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 = int (*)();
    using CreateFunction = int (*)(const char*, yolo_runtime_t**);
    using GetCapabilitiesFunction = unsigned int (*)(const yolo_runtime_t*);
    using InferImageFunction = int (*)(yolo_runtime_t*, const yolo_image_view_t*,
                                       yolo_result_list_t*);
    using LastErrorFunction = const char* (*)(const yolo_runtime_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_t image{};
        image.struct_size = sizeof(image);
        image.abi_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_result_list_t output{};
        const int status = m_inferImage(m_runtime, &image, &output);
        if (status != YOLO_STATUS_OK) {
            const QString runtimeError = RuntimeError();
            if (status == YOLO_STATUS_INITIALIZATION_FAILED ||
                status == YOLO_STATUS_BACKEND_FAILED) {
                UnloadRuntimeLibrary();
            }
            errorMessage = QStringLiteral("AAPGS模型推理失败(%1)：%2")
                               .arg(status)
                               .arg(runtimeError);
            return false;
        }
        if (output.task != YOLO_TASK_CLS) {
            errorMessage = QStringLiteral("AAPGS模型任务不是分类任务：%1")
                               .arg(static_cast<int>(output.task));
            return false;
        }
        if (output.count <= 0 || output.count > YOLO_MAX_RESULTS) {
            errorMessage = QStringLiteral("AAPGS分类结果数量异常：%1")
                               .arg(output.count);
            return false;
        }

        const yolo_result_t* bestResult = nullptr;
        for (int index = 0; index < output.count; ++index) {
            const yolo_result_t& item = output.results[index];
            if (!std::isfinite(item.score) || item.score < 0.0f ||
                item.score > 1.0f) {
                continue;
            }
            if (!bestResult || item.score > bestResult->score) {
                bestResult = &item;
            }
        }
        if (!bestResult) {
            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;
        }

#if !defined(Q_OS_LINUX) || !defined(Q_PROCESSOR_ARM_64)
        errorMessage = QStringLiteral(
            "AAPGS机型分类仅支持RK3588 Linux AArch64");
        return false;
#else
        QString configPath;
        QString packageRoot;
        if (!FindModelConfig(configPath, packageRoot, errorMessage) ||
            !LoadRuntimeLibrary(packageRoot, errorMessage)) {
            return false;
        }

        const int 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);
        const int status = m_create(encodedConfigPath.constData(), &m_runtime);
        if (status != YOLO_STATUS_OK || !m_runtime) {
            const QString runtimeError = RuntimeError();
            UnloadRuntimeLibrary();
            errorMessage = QStringLiteral("AAPGS模型初始化失败(%1)：%2；配置：%3")
                               .arg(status)
                               .arg(runtimeError)
                               .arg(configPath);
            return false;
        }
        if ((m_getCapabilities(m_runtime) & YOLO_CAPABILITY_INFER_IMAGE) == 0) {
            UnloadRuntimeLibrary();
            errorMessage = QStringLiteral("AAPGS运行库不支持内存图像推理");
            return false;
        }
        return true;
#endif
    }

    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);
        }
        AppendUniquePath(candidates, QDir(packageRoot).filePath(
            QStringLiteral("lib/libyolo_runtime.so.3")));
        AppendUniquePath(candidates, QDir(packageRoot).filePath(
            QStringLiteral("lib/libyolo_runtime.so")));
        AppendUniquePath(candidates, QStringLiteral("libyolo_runtime.so.3"));
        AppendUniquePath(candidates, QStringLiteral("yolo_runtime"));

        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_getCapabilities = reinterpret_cast<GetCapabilitiesFunction>(
            m_library.resolve("yolo_runtime_get_capabilities"));
        m_inferImage = reinterpret_cast<InferImageFunction>(
            m_library.resolve("yolo_runtime_infer_image"));
        m_lastError = reinterpret_cast<LastErrorFunction>(
            m_library.resolve("yolo_runtime_last_error"));
        m_destroy = reinterpret_cast<DestroyFunction>(
            m_library.resolve("yolo_runtime_destroy"));
        return m_getAbiVersion && m_create && m_getCapabilities &&
               m_inferImage && m_lastError && m_destroy;
    }

    QString RuntimeError() const
    {
        if (!m_runtime || !m_lastError) {
            return QStringLiteral("未知错误");
        }
        const char* error = m_lastError(m_runtime);
        const QString text = error ? QString::fromUtf8(error).trimmed() : QString();
        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_getCapabilities = nullptr;
        m_inferImage = nullptr;
        m_lastError = 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;
    GetCapabilitiesFunction m_getCapabilities = nullptr;
    InferImageFunction m_inferImage = nullptr;
    LastErrorFunction m_lastError = 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);
}