ViewDesignEngine/include/vde/ai/inference_engine.h

#pragma once
/**
 * @file inference_engine.h
 * @brief AI 推理引擎：ONNX Runtime 集成 + TensorRT 加速 + 模型版本管理
 *
 * ## 主要功能
 *
 * | 功能                  | 说明                                          |
 * |-----------------------|----------------------------------------------|
 * | InferenceSession      | ONNX Runtime 推理会话封装                     |
 * | TensorRTAccelerator   | TensorRT 构建/优化/推理加速                   |
 * | ModelRegistry         | 模型版本管理与自动选择                        |
 *
 * ## 使用模式
 *
 * @code{.cpp}
 * vde::ai::InferenceSession session;
 * session.load_model("model.onnx");
 * auto output = session.run(input_tensor);
 *
 * // TensorRT 加速（需 NVIDIA GPU + CUDA）
 * vde::ai::TensorRTAccelerator trt;
 * trt.build_engine("model.onnx");
 * auto fast_output = trt.run(input_tensor);
 *
 * // 模型版本管理
 * vde::ai::ModelRegistry registry;
 * registry.register_model("classifier", "v2.0", "model_v2.onnx");
 * auto best = registry.get_best_model("classifier");  // 自动选择最新兼容版本
 * @endcode
 *
 * @ingroup ai
 */
#include <vector>
#include <string>
#include <memory>
#include <unordered_map>
#include <optional>
#include <functional>
#include <cstdint>

namespace vde::ai {

// ═══════════════════════════════════════════════════════════
//  通用张量类型
// ═══════════════════════════════════════════════════════════

/** @brief 数据类型枚举 */
enum class TensorDataType : int {
    Float32 = 0,
    Float64 = 1,
    Int32   = 3,
    Int64   = 7,
};

/** @brief 张量容器 */
struct Tensor {
    std::string name;                   ///< 张量名称（对应 ONNX 节点名）
    std::vector<int64_t> shape;         ///< 形状 (N, C, H, W, ...)
    TensorDataType dtype = TensorDataType::Float32;
    std::vector<uint8_t> data;          ///< 原始二进制数据
    size_t element_count = 0;           ///< 元素总数

    /** @brief 获取数据的 float* 视图 */
    [[nodiscard]] const float* data_f32() const;
    /** @brief 获取数据的 float* 可写视图 */
    [[nodiscard]] float* mutable_data_f32();
    /** @brief 获取数据的 int64_t* 视图 */
    [[nodiscard]] const int64_t* data_i64() const;
    /** @brief 获取字节大小 */
    [[nodiscard]] size_t byte_size() const;

    /** @brief 创建给定形状的 float32 张量 */
    static Tensor make_f32(const std::vector<int64_t>& shape);
    /** @brief 创建给定形状的 int64 张量 */
    static Tensor make_i64(const std::vector<int64_t>& shape);
};

// ═══════════════════════════════════════════════════════════
//  Provider 类型
// ═══════════════════════════════════════════════════════════

/** @brief ONNX Runtime 执行提供者 */
enum class ExecutionProvider {
    CPU,                ///< CPU 执行（默认，总是可用）
    CUDA,               ///< CUDA GPU 加速
    TensorRT,           ///< NVIDIA TensorRT
    OpenVINO,           ///< Intel OpenVINO
    DirectML            ///< Microsoft DirectML (Windows)
};

/** @brief 设备信息 */
struct DeviceInfo {
    ExecutionProvider provider = ExecutionProvider::CPU;
    int device_id = 0;              ///< 设备编号
    std::string device_name;        ///< 设备名称（如 "NVIDIA GeForce RTX 4090"）
    int64_t memory_mb = 0;          ///< 可用显存/内存 (MB)
    bool available = false;         ///< 是否可用
};

// ═══════════════════════════════════════════════════════════
//  InferenceSession — ONNX Runtime 推理会话
// ═══════════════════════════════════════════════════════════

/** @brief 会话配置 */
struct InferenceConfig {
    ExecutionProvider preferred_provider = ExecutionProvider::CPU;
    int intra_op_threads = 4;           ///< 算子内并行线程数
    int inter_op_threads = 2;           ///< 算子间并行线程数
    bool enable_profiling = false;      ///< 是否开启性能分析
    int graph_optimization_level = 2;   ///< 图优化级别 (0=无, 1=基础, 2=扩展, 99=ALL)
    bool enable_memory_pattern = true;  ///< 是否启用内存重用模式
    int log_severity_level = 2;         ///< 日志级别 (0=VERBOSE, 1=INFO, 2=WARNING, 3=ERROR, 4=FATAL)
};

/** @brief 推理统计信息 */
struct InferenceStats {
    double inference_time_ms = 0.0;     ///< 本轮推理耗时
    double avg_inference_time_ms = 0.0; ///< 平均推理耗时
    int64_t total_inferences = 0;       ///< 总推理次数
    size_t input_size_bytes = 0;        ///< 输入大小
    size_t output_size_bytes = 0;       ///< 输出大小
};

/**
 * @brief ONNX Runtime 推理会话封装
 *
 * 管理 ONNX 模型加载、输入/输出绑定和推理执行。
 */
class InferenceSession {
public:
    InferenceSession();
    explicit InferenceSession(const InferenceConfig& config);
    ~InferenceSession();

    // ── 模型加载 ──────────────────────────────────

    /**
     * @brief 加载 ONNX 模型
     * @param model_path 模型文件路径 (.onnx)
     * @return 是否成功
     */
    bool load_model(const std::string& model_path);

    /**
     * @brief 从内存加载 ONNX 模型
     * @param model_data 模型字节数据
     * @param data_size 数据大小
     * @return 是否成功
     */
    bool load_model_from_memory(const uint8_t* model_data, size_t data_size);

    /** @brief 模型是否已加载 */
    [[nodiscard]] bool is_loaded() const noexcept;

    // ── 模型信息 ──────────────────────────────────

    /** @brief 获取输入张量名称和形状 */
    [[nodiscard]] std::unordered_map<std::string, std::vector<int64_t>> input_info() const;

    /** @brief 获取输出张量名称和形状 */
    [[nodiscard]] std::unordered_map<std::string, std::vector<int64_t>> output_info() const;

    /** @brief 获取模型元数据（作者、版本等） */
    [[nodiscard]] std::unordered_map<std::string, std::string> model_metadata() const;

    // ── 推理 ──────────────────────────────────────

    /**
     * @brief 执行推理
     * @param inputs 输入张量列表
     * @return 输出张量列表
     */
    [[nodiscard]] std::vector<Tensor> run(const std::vector<Tensor>& inputs);

    /**
     * @brief 执行推理（单输入单输出便捷接口）
     * @param input 输入张量
     * @return 输出张量
     */
    [[nodiscard]] Tensor run_single(const Tensor& input);

    // ── 性能 ──────────────────────────────────────

    /** @brief 获取最新推理统计 */
    [[nodiscard]] const InferenceStats& stats() const noexcept { return stats_; }

    /** @brief 重置统计 */
    void reset_stats();

    /** @brief 获取可用执行提供者列表 */
    [[nodiscard]] std::vector<DeviceInfo> available_providers() const;

    // ── 配置 ──────────────────────────────────────

    /** @brief 获取当前配置 */
    [[nodiscard]] const InferenceConfig& config() const noexcept { return config_; }

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
    InferenceConfig config_;
    InferenceStats stats_;
};

// ═══════════════════════════════════════════════════════════
//  TensorRTAccelerator — TensorRT 加速
// ═══════════════════════════════════════════════════════════

/** @brief TensorRT 引擎构建配置 */
struct TensorRTConfig {
    std::string onnx_model_path;        ///< 源 ONNX 模型路径
    std::string engine_cache_path;      ///< TensorRT 引擎缓存路径
    int max_batch_size = 1;             ///< 最大批大小
    int max_workspace_size_mb = 4096;   ///< 最大工作空间 (MB)
    bool use_fp16 = false;              ///< 是否使用 FP16 半精度
    bool use_int8 = false;              ///< 是否使用 INT8 量化
    bool strict_type_constraints = true;///< 严格类型约束
    int device_id = 0;                  ///< CUDA 设备编号
};

/** @brief TensorRT 加速器状态 */
struct TensorRTStatus {
    bool engine_built = false;          ///< 引擎是否已构建
    bool engine_loaded = false;         ///< 引擎是否已加载
    double build_time_ms = 0.0;         ///< 构建耗时
    double inference_time_ms = 0.0;     ///< 最近一次推理耗时
    int64_t total_inferences = 0;       ///< 总推理次数
    int64_t engine_size_bytes = 0;      ///< 序列化引擎大小
};

/**
 * @brief TensorRT 推理加速器
 *
 * 将 ONNX 模型编译为 TensorRT 引擎以实现高性能推理。
 * 需要 NVIDIA GPU + CUDA + TensorRT 库。
 */
class TensorRTAccelerator {
public:
    TensorRTAccelerator();
    explicit TensorRTAccelerator(const TensorRTConfig& config);
    ~TensorRTAccelerator();

    // ── 引擎构建 ──────────────────────────────────

    /**
     * @brief 从 ONNX 模型构建 TensorRT 引擎
     * @param onnx_model_path ONNX 模型路径
     * @param engine_cache_path 引擎缓存保存路径（空字符串=不缓存）
     * @return 是否成功
     */
    bool build_engine(const std::string& onnx_model_path,
                      const std::string& engine_cache_path = "");

    /**
     * @brief 从缓存加载已编译的 TensorRT 引擎
     * @param engine_path 引擎文件路径
     * @return 是否成功
     */
    bool load_engine(const std::string& engine_path);

    // ── 推理 ──────────────────────────────────────

    /**
     * @brief 执行 TensorRT 推理
     * @param inputs 输入张量
     * @return 输出张量
     */
    [[nodiscard]] std::vector<Tensor> run(const std::vector<Tensor>& inputs);

    // ── 状态 ──────────────────────────────────────

    /** @brief 检查是否可用（CUDA + TensorRT） */
    [[nodiscard]] static bool is_available() noexcept;

    /** @brief 获取状态 */
    [[nodiscard]] const TensorRTStatus& status() const noexcept { return status_; }

    /** @brief 获取配置 */
    [[nodiscard]] const TensorRTConfig& config() const noexcept { return config_; }

    // ── 性能 ──────────────────────────────────────

    /**
     * @brief 预热引擎（执行若干次空推理以稳定 GPU 频率）
     * @param warmup_runs 预热运行次数
     */
    void warmup(int warmup_runs = 10);

    /**
     * @brief 性能基准测试
     * @param input_shape 输入张量形状
     * @param iterations 迭代次数
     * @return 平均推理耗时 (ms)
     */
    [[nodiscard]] double benchmark(const std::vector<int64_t>& input_shape, int iterations = 100);

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
    TensorRTConfig config_;
    TensorRTStatus status_;
};

// ═══════════════════════════════════════════════════════════
//  ModelRegistry — 模型版本管理
// ═══════════════════════════════════════════════════════════

/** @brief 模型条目 */
struct ModelEntry {
    std::string name;               ///< 模型名称（如 "feature_classifier"）
    std::string version;            ///< 语义版本号（如 "1.0.0"）
    std::string path;               ///< 模型文件路径
    std::string checksum;           ///< 文件 SHA256 校验和
    std::string framework;          ///< 框架来源（"onnx", "tensorrt", "custom"）
    std::unordered_map<std::string, std::string> metadata; ///< 附加元数据
    int64_t file_size = 0;          ///< 文件大小（字节）
    int64_t registered_at = 0;      ///< 注册时间戳
    bool is_active = true;          ///< 是否活跃
};

/** @brief 模型注册表配置 */
struct ModelRegistryConfig {
    std::string registry_root;      ///< 模型存储根目录
    std::string default_provider;   ///< 默认执行提供者
    int max_versions_per_model = 5; ///< 每个模型保留的最大版本数
    bool auto_cleanup = true;       ///< 是否自动清理旧版本
};

/**
 * @brief AI 模型版本管理器
 *
 * 管理多个模型的版本注册、查找和生命周期。
 *
 * ## 功能
 *
 * - 注册模型版本（名称 + 语义版本号 + 文件路径）
 * - 获取最新版本、特定版本或满足约束的最佳版本
 * - 自动清理过时版本
 * - 校验和验证
 */
class ModelRegistry {
public:
    explicit ModelRegistry(const ModelRegistryConfig& config = ModelRegistryConfig{});
    ~ModelRegistry();

    // ── 注册 ──────────────────────────────────────

    /**
     * @brief 注册一个模型版本
     * @param entry 模型条目
     * @return 是否成功
     */
    bool register_model(const ModelEntry& entry);

    /**
     * @brief 便捷注册
     * @param name 模型名称
     * @param version 版本号
     * @param path 文件路径
     * @param framework 框架类型
     * @param metadata 附加元数据
     * @return 是否成功
     */
    bool register_model(const std::string& name,
                        const std::string& version,
                        const std::string& path,
                        const std::string& framework = "onnx",
                        const std::unordered_map<std::string, std::string>& metadata = {});

    // ── 查询 ──────────────────────────────────────

    /**
     * @brief 获取指定模型的最新活跃版本
     * @param name 模型名称
     * @return 模型条目（如不存在返回 nullopt）
     */
    [[nodiscard]] std::optional<ModelEntry> get_latest(const std::string& name) const;

    /**
     * @brief 获取特定版本
     * @param name 模型名称
     * @param version 版本号
     * @return 模型条目（如不存在返回 nullopt）
     */
    [[nodiscard]] std::optional<ModelEntry> get_version(
        const std::string& name, const std::string& version) const;

    /**
     * @brief 获取满足要求的最佳版本（选择最新兼容版本）
     * @param name 模型名称
     * @param min_version 最低版本要求（如 "1.2.0"）
     * @return 最佳兼容版本
     */
    [[nodiscard]] std::optional<ModelEntry> get_best_model(
        const std::string& name, const std::string& min_version = "0.0.0") const;

    /**
     * @brief 列出某模型的所有已注册版本
     * @param name 模型名称
     * @return 版本列表（按版本号降序）
     */
    [[nodiscard]] std::vector<ModelEntry> list_versions(const std::string& name) const;

    /**
     * @brief 列出所有已注册模型名称
     * @return 模型名称列表
     */
    [[nodiscard]] std::vector<std::string> list_models() const;

    // ── 管理 ──────────────────────────────────────

    /**
     * @brief 停用某个版本
     * @param name 模型名称
     * @param version 版本号
     * @return 是否成功
     */
    bool deactivate(const std::string& name, const std::string& version);

    /**
     * @brief 删除某个版本
     * @param name 模型名称
     * @param version 版本号
     * @return 是否成功
     */
    bool remove(const std::string& name, const std::string& version);

    /**
     * @brief 清理过时版本（保留最新的 max_versions_per_model 个）
     */
    void cleanup();

    /**
     * @brief 验证所有模型文件的校验和
     * @return (valid_count, invalid_count) pair
     */
    [[nodiscard]] std::pair<int, int> verify_checksums() const;

    /** @brief 获取已注册模型总数 */
    [[nodiscard]] size_t total_entries() const noexcept;

    /** @brief 获取配置 */
    [[nodiscard]] const ModelRegistryConfig& config() const noexcept { return config_; }

    // ── 静态工具 ──────────────────────────────────

    /**
     * @brief 比较两个语义版本号
     * @return -1 (a<b), 0 (a==b), 1 (a>b)
     */
    [[nodiscard]] static int compare_versions(const std::string& a, const std::string& b);

    /**
     * @brief 计算文件的 SHA256 校验和
     * @param filepath 文件路径
     * @return 十六进制校验和字符串
     */
    [[nodiscard]] static std::string compute_sha256(const std::string& filepath);

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
    ModelRegistryConfig config_;
};

} // namespace vde::ai