"""
检查环境依赖 - 简化版
"""

import sys


def check_package(package_name, min_version=None):
    """检查包是否安装及版本"""
    try:
        import importlib
        module = importlib.import_module(package_name)
        
        if hasattr(module, "__version__"):
            installed_version = module.__version__
            status = f"{package_name}: {installed_version}"
        else:
            status = f"{package_name}: 已安装"
        
        print(f"[OK] {status}")
        return True
    except ImportError:
        print(f"[MISSING] {package_name}")
        return False


def check_cuda():
    """检查 CUDA 是否可用"""
    try:
        import torch
        if torch.cuda.is_available():
            print(f"[OK] CUDA: 可用 ({torch.cuda.device_count()} 个 GPU)")
            for i in range(torch.cuda.device_count()):
                print(f"     GPU {i}: {torch.cuda.get_device_name(i)}")
            return True
        else:
            print(f"[WARN] CUDA: 不可用 (将使用 CPU 训练)")
            return False
    except ImportError:
        print(f"[MISSING] PyTorch")
        return False


def main():
    print("=" * 60)
    print("FineTuneX 环境检查")
    print("=" * 60)
    print()
    
    required_packages = [
        "torch",
        "transformers",
        "datasets",
        "accelerate",
        "peft",
        "bitsandbytes",
        "trl",
        "fastapi",
        "uvicorn",
        "pydantic",
    ]
    
    all_ok = True
    
    print("检查 Python 包:")
    print("-" * 60)
    for package in required_packages:
        if not check_package(package):
            all_ok = False
    
    print()
    print("检查 CUDA:")
    print("-" * 60)
    if not check_cuda():
        print("提示：CPU 训练速度较慢，建议使用 GPU")
    
    print()
    print("=" * 60)
    if all_ok:
        print("[OK] 所有核心依赖已安装！")
    else:
        print("[ERROR] 部分依赖缺失")
        print("请运行：pip install -r requirements.txt")
    print("=" * 60)
    
    return 0 if all_ok else 1


if __name__ == "__main__":
    sys.exit(main())