#!/usr/bin/env python
"""环境自检：确认依赖与可用的计算后端。

    python scripts/check_env.py
"""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from utils.device import (  # noqa: E402
    available_backends,
    default_amp,
    describe_device,
    memory_stats,
    resolve_device,
)


def main() -> int:
    print(f"Python: {sys.version.split()[0]}")
    try:
        import torch
    except ImportError:
        print("PyTorch: 未安装 —— 先执行 `uv sync`")
        return 1

    print(f"PyTorch: {torch.__version__}  (CUDA 构建: {torch.version.cuda})")
    print(f"可用后端（按优先级）: {available_backends()}")

    print("\n后端明细:")
    print(f"  cuda: available={torch.cuda.is_available()}  device_count={torch.cuda.device_count()}")
    print(
        f"  mps : available={torch.backends.mps.is_available()}  "
        f"built={torch.backends.mps.is_built()}"
    )

    device = resolve_device("auto")
    print(f"\ndevice=auto 会选中: {describe_device(device)}")
    print(f"  AMP 默认: {'开' if default_amp(device) else '关'}")
    stats = memory_stats(device)
    if stats:
        print("  " + "  ".join(f"{k}={v:.0f}" for k, v in stats.items()))

    if device.type == "cpu":
        print("\n⚠️  只有 CPU 可用，训练会非常慢（实测比 MPS 慢约 30 倍），建议换有 GPU 的机器。")
        return 1
    print("\n训练环境就绪。")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())