#!/usr/bin/env python3
"""
setup_check.py - Verify all dependencies are installed correctly

Run this after installing dependencies to ensure everything works.
"""

import sys
import os

def check_imports():
    """Check all required imports work."""
    print("Checking Python packages...")
    
    packages = [
        ("customtkinter", "GUI Framework"),
        ("numpy", "Numerical Computing"),
        ("scipy", "Signal Processing"),
        ("soundfile", "Audio I/O"),
        ("pygame", "Audio Playback"),
        ("pyloudnorm", "Loudness Measurement"),
    ]
    
    optional_packages = [
        ("pedalboard", "Audio Effects (Spotify)"),
        ("torch", "PyTorch (Deep Learning)"),
        ("torchaudio", "TorchAudio (Audio ML)"),
        ("demucs", "Stem Separation (Meta)"),
    ]
    
    all_ok = True
    
    print("\n[Required Packages]")
    for package, description in packages:
        try:
            __import__(package)
            print(f"  ✓ {package} - {description}")
        except ImportError as e:
            print(f"  ✗ {package} - {description}: {e}")
            all_ok = False
    
    print("\n[Optional Packages (for stem separation)]")
    for package, description in optional_packages:
        try:
            __import__(package)
            print(f"  ✓ {package} - {description}")
        except ImportError as e:
            print(f"  ⚠ {package} - {description}: Not installed")
    
    return all_ok


def check_cuda():
    """Check CUDA availability for GPU acceleration."""
    print("\n[GPU Acceleration]")
    
    try:
        import torch
        if torch.cuda.is_available():
            print(f"  ✓ CUDA available: {torch.cuda.get_device_name(0)}")
            print(f"    - CUDA version: {torch.version.cuda}")
            print(f"    - GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
        else:
            print("  ⚠ CUDA not available - will use CPU (slower stem separation)")
    except ImportError:
        print("  ⚠ PyTorch not installed - cannot check GPU")


def check_audio_device():
    """Check audio output device."""
    print("\n[Audio Output]")
    
    try:
        import pygame
        pygame.mixer.init(frequency=44100, size=-16, channels=2, buffer=2048)
        print(f"  ✓ Audio initialized successfully")
        pygame.mixer.quit()
    except Exception as e:
        print(f"  ✗ Audio initialization failed: {e}")


def main():
    print("=" * 60)
    print("SunoRemaster - Dependency Check")
    print("=" * 60)
    
    packages_ok = check_imports()
    check_cuda()
    check_audio_device()
    
    print("\n" + "=" * 60)
    if packages_ok:
        print("✓ All required packages installed!")
        print("  Run: ./run.sh to start the application")
    else:
        print("✗ Some packages are missing.")
        print("  Run: ./venv310/bin/pip install -r requirements.txt")
    print("=" * 60)
    
    return 0 if packages_ok else 1


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