"""
TensorRT 环境诊断脚本
"""
import os
import sys

def check_env():
    print("=" * 60)
    print("TensorRT 环境诊断")
    print("=" * 60)
    
    # 1. Python 版本
    print(f"\n[1] Python: {sys.version}")
    
    # 2. CUDA 环境
    print("\n[2] CUDA 环境:")
    cuda_path = os.environ.get("CUDA_PATH", "未设置")
    print(f"    CUDA_PATH: {cuda_path}")
    
    # 3. PyTorch CUDA
    print("\n[3] PyTorch:")
    try:
        import torch
        print(f"    版本: {torch.__version__}")
        print(f"    CUDA 可用: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"    CUDA 版本: {torch.version.cuda}")
            print(f"    GPU: {torch.cuda.get_device_name(0)}")
    except ImportError:
        print("    未安装")
    
    # 4. TensorRT Python 包
    print("\n[4] TensorRT Python 包:")
    try:
        import tensorrt as trt
        print(f"    版本: {trt.__version__}")
        print(f"    文件位置: {trt.__file__}")
    except ImportError as e:
        print(f"    导入失败: {e}")
    
    # 5. PyCUDA
    print("\n[5] PyCUDA:")
    try:
        import pycuda.driver as cuda
        import pycuda.autoinit
        print(f"    版本: {cuda.get_version()}")
    except ImportError as e:
        print(f"    导入失败: {e}")
    except Exception as e:
        print(f"    初始化失败: {e}")
    
    # 6. Ultralytics
    print("\n[6] Ultralytics:")
    try:
        import ultralytics
        print(f"    版本: {ultralytics.__version__}")
    except ImportError:
        print("    未安装")
    
    # 7. 测试 TensorRT Runtime 加载
    print("\n[7] TensorRT Runtime 测试:")
    try:
        import tensorrt as trt
        logger = trt.Logger(trt.Logger.WARNING)
        runtime = trt.Runtime(logger)
        print(f"    Runtime 创建成功")
        print(f"    TRT 序列化版本: 检查 engine 文件...")
    except Exception as e:
        print(f"    失败: {e}")
    
    # 8. 检查 engine 文件
    print("\n[8] Engine 文件检查:")
    engine_paths = [
        "C:/Users/16337/PycharmProjects/Security_project/yolov8n.engine",
        "C:/Users/16337/PycharmProjects/Security_project/yolov8n_320.engine",
    ]
    for path in engine_paths:
        if os.path.exists(path):
            size = os.path.getsize(path) / (1024 * 1024)
            print(f"    {path}")
            print(f"      大小: {size:.2f} MB")
            # 读取文件头
            with open(path, 'rb') as f:
                header = f.read(32)
                print(f"      文件头 (hex): {header[:16].hex()}")
        else:
            print(f"    {path} - 不存在")
    
    # 9. 尝试加载 engine
    print("\n[9] 尝试加载 Engine:")
    for path in engine_paths:
        if os.path.exists(path):
            try:
                import tensorrt as trt
                logger = trt.Logger(trt.Logger.ERROR)
                runtime = trt.Runtime(logger)
                with open(path, 'rb') as f:
                    engine = runtime.deserialize_cuda_engine(f.read())
                if engine:
                    print(f"    {os.path.basename(path)}: 加载成功 ✓")
                else:
                    print(f"    {os.path.basename(path)}: 加载失败 (engine=None)")
            except Exception as e:
                print(f"    {os.path.basename(path)}: 加载失败 - {e}")
    
    print("\n" + "=" * 60)
    print("诊断完成")
    print("=" * 60)

if __name__ == "__main__":
    check_env()