"""
install.py — Detect GPU and install the right packages automatically.
Run once before first use: python install.py
"""

import subprocess
import sys
import platform
import shutil

def run(cmd, check=True):
    print(f"  $ {' '.join(cmd)}")
    result = subprocess.run(cmd, check=check)
    return result.returncode == 0

def detect_cuda():
    """Check if nvidia-smi is available (implies CUDA-capable GPU)."""
    return shutil.which("nvidia-smi") is not None

def detect_rocm():
    """Check for AMD ROCm (Linux only)."""
    return shutil.which("rocm-smi") is not None

def main():
    print("=" * 60)
    print("SubtitleBurner — Installer")
    print("=" * 60)
    print(f"Python: {sys.version}")
    print(f"Platform: {platform.system()} {platform.machine()}")
    print()

    pip = [sys.executable, "-m", "pip", "install", "--upgrade"]

    # Core deps always needed
    print("Installing core dependencies…")
    run(pip + ["psutil"])

    # faster-whisper
    print("\nInstalling faster-whisper…")
    run(pip + ["faster-whisper"])

    # PyTorch — pick the right variant
    has_cuda = detect_cuda()
    has_rocm = detect_rocm()
    is_mac = platform.system() == "Darwin"

    if has_cuda:
        print("\n✅ NVIDIA GPU detected — installing PyTorch with CUDA 12.1 support…")
        run(pip + [
            "torch", "torchvision",
            "--index-url", "https://download.pytorch.org/whl/cu121"
        ])
    elif has_rocm and platform.system() == "Linux":
        print("\n✅ AMD GPU detected — installing PyTorch with ROCm 6.0 support…")
        run(pip + [
            "torch", "torchvision",
            "--index-url", "https://download.pytorch.org/whl/rocm6.0"
        ])
    elif is_mac:
        print("\nApple Silicon / Mac — installing standard PyTorch (MPS support built in)…")
        run(pip + ["torch", "torchvision"])
    else:
        print("\nNo GPU detected — installing CPU-only PyTorch…")
        run(pip + [
            "torch", "torchvision",
            "--index-url", "https://download.pytorch.org/whl/cpu"
        ])

    # Check ffmpeg
    print()
    if shutil.which("ffmpeg"):
        print("✅ ffmpeg found in PATH.")
    else:
        print("⚠  ffmpeg NOT found.")
        if platform.system() == "Windows":
            print("   → Download from https://ffmpeg.org/download.html")
            print("   → Or run: winget install ffmpeg")
        elif platform.system() == "Darwin":
            print("   → Run: brew install ffmpeg")
        else:
            print("   → Run: sudo apt install ffmpeg  (or your distro's package manager)")

    print()
    print("=" * 60)
    print("Installation complete! Run the app with:")
    print("  python subtitle_burner.py")
    print("=" * 60)


if __name__ == "__main__":
    main()