#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.yaml";
const char kCanonicalModelConfigRelativePath[] =
    "models/model_20260712T020258_0000_6339d04b/"
    "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 < 4; ++level) {
        if (QFileInfo::exists(directory.filePath(
                QStringLiteral("lib/libyolo_runtime.so.2")))) {
            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()) {
            configuredFile.setFile(QDir(configuredFile.absoluteFilePath())
                                       .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"));
    AppendUniquePath(roots, QStringLiteral("/usr/lib/AAPGS_model"));
    AppendUniquePath(roots, QStringLiteral("/usr/local/lib/AAPGS_model"));

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

    errorMessage = QStringLiteral(
        "未找到AAPGS模型配置；DEB应安装到"
        "/opt/rk3588-ai/models/model.yaml，或设置AAPGS_MODEL_CONFIG");
    return false;
}

QString ModelName(const yolo_result_t& item, QString& errorMessage)
{
    errorMessage.clear();
    int labelLength = 0;
    while (labelLength < YOLO_LABEL_LENGTH && item.label[labelLength] != '\0') {
        ++labelLength;
    }

    const QString label = QString::fromUtf8(item.label, labelLength).trimmed();
    if (label.compare(QStringLiteral("a320"), Qt::CaseInsensitive) == 0) {
        if (item.class_id != 0) {
            errorMessage = QStringLiteral("AAPGS类别标签与编号不一致：%1/%2")
                .arg(label)
                .arg(item.class_id);
            return QString();
        }
        return QStringLiteral("A320");
    }
    if (label.compare(QStringLiteral("b737"), Qt::CaseInsensitive) == 0) {
        if (item.class_id != 1) {
            errorMessage = QStringLiteral("AAPGS类别标签与编号不一致：%1/%2")
                .arg(label)
                .arg(item.class_id);
            return QString();
        }
        return QStringLiteral("B737");
    }

    if (label.isEmpty()) {
        if (item.class_id == 0) {
            return QStringLiteral("A320");
        }
        if (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("机型识别图像无效");
            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("机型识别图像转换失败");
            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& candidate = output.results[index];
            if (!std::isfinite(candidate.score) ||
                candidate.score < 0.0f || candidate.score > 1.0f) {
                continue;
            }
            if (!bestResult || candidate.score > bestResult->score) {
                bestResult = &candidate;
            }
        }
        if (!bestResult) {
            errorMessage = QStringLiteral("AAPGS未返回有效分类结果");
            return false;
        }

        QString classificationError;
        result.modelType = ModelName(*bestResult, classificationError);
        result.confidence = static_cast<double>(bestResult->score);
        if (result.modelType.isEmpty()) {
            errorMessage = classificationError.trimmed().isEmpty()
                ? QStringLiteral("AAPGS返回了无效机型类别")
                : classificationError;
            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)) {
            return false;
        }
        if (!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.2")));
        AppendUniquePath(candidates, QDir(packageRoot).filePath(
            QStringLiteral("lib/libyolo_runtime.so")));
        AppendUniquePath(candidates, QStringLiteral("libyolo_runtime.so.2"));
        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 ResetRuntime()
    {
        if (m_runtime && m_destroy) {
            m_destroy(m_runtime);
        }
        m_runtime = nullptr;
    }

    void ClearFunctions()
    {
        m_getAbiVersion = nullptr;
        m_create = nullptr;
        m_getCapabilities = nullptr;
        m_inferImage = nullptr;
        m_lastError = nullptr;
        m_destroy = nullptr;
    }

    void UnloadRuntimeLibrary()
    {
        ResetRuntime();
        if (m_library.isLoaded()) {
            m_library.unload();
        }
        ClearFunctions();
    }

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);
}