using Check.Main.Common;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using YoloDotNet;
using YoloDotNet.Models;

namespace Check.Main.Infer
{
    /// <summary>
    /// 静态全局YOLO模型管理器。
    /// 负责在程序启动时加载所有模型，并在关闭时释放资源。
    /// </summary>
    public static class YoloModelManager
    {
        // 使用 ConcurrentDictionary 保证线程安全
        private static readonly ConcurrentDictionary<int, Yolo> _loadedModels = new ConcurrentDictionary<int, Yolo>();

        /// <summary>
        /// 根据模型配置列表初始化所有YOLO模型实例。
        /// 应在程序启动时调用一次。
        /// </summary>
        /// <param name="modelSettings">模型配置列表。</param>
        public static void Initialize(List<ModelSettings> modelSettings)
        {
            Shutdown(); // 先确保清理掉旧实例

            if (modelSettings == null) return;

            ThreadSafeLogger.Log("开始加载YOLO模型...");
            foreach (var setting in modelSettings)
            {
                bool gpuUse = false;
                if (setting.CheckDevice == DetectDevice.GPU)
                {
                    gpuUse = true;
                }

                if (string.IsNullOrEmpty(setting.Path) || !File.Exists(setting.Path))
                {
                    ThreadSafeLogger.Log($"[警告] 模型 '{setting.Name}' (ID: {setting.Id}) 路径无效或文件不存在，已跳过加载。");
                    continue;
                }

                try
                {
                    // 创建YOLO实例
                    var yolo = new Yolo(new YoloOptions
                    {
                        OnnxModel = setting.Path,
                        // 您可以根据需要从配置中读取这些值
                        ModelType = (YoloDotNet.Enums.ModelType)setting.MType,
                        Cuda = gpuUse, // 推荐使用GPU
                        PrimeGpu = false
                    });

                    if (_loadedModels.TryAdd(setting.Id, yolo))
                    {
                        ThreadSafeLogger.Log($"成功加载模型 '{setting.Name}' (ID: {setting.Id})。");
                    }
                }
                catch (Exception ex)
                {
                    ThreadSafeLogger.Log($"[错误] 加载模型 '{setting.Name}' (ID: {setting.Id}) 失败: {ex.Message}");
                }
            }
            ThreadSafeLogger.Log($"YOLO模型加载完成，共成功加载 {_loadedModels.Count} 个模型。");
        }

        /// <summary>
        /// 获取指定ID的已加载YOLO模型。
        /// </summary>
        /// <param name="modelId">模型编号。</param>
        /// <returns>Yolo实例，如果未找到则返回null。</returns>
        public static Yolo GetModel(int modelId)
        {
            _loadedModels.TryGetValue(modelId, out var yolo);
            return yolo;
        }

        /// <summary>
        /// 释放所有已加载的YOLO模型资源。
        /// 应在程序关闭时调用。
        /// </summary>
        public static void Shutdown()
        {
            if (_loadedModels.Count > 0)
            {
                ThreadSafeLogger.Log("正在释放所有YOLO模型...");
                foreach (var yolo in _loadedModels.Values)
                {
                    yolo?.Dispose();
                }
                _loadedModels.Clear();
                ThreadSafeLogger.Log("所有YOLO模型已释放。");
            }
        }
    }
}