"""
Test your setup locally before deploying to HF
This helps identify issues without waiting for HF Space builds
"""
import os
import sys
import torch

print("=" * 80)
print(" " * 25 + "Local Setup Test")
print("=" * 80)

# Test 1: Python version
print("\n[1/7] Python Version")
print(f"✓ Python {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}")

# Test 2: PyTorch
print("\n[2/7] PyTorch")
try:
    print(f"✓ PyTorch version: {torch.__version__}")
    print(f"✓ CUDA available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"✓ CUDA version: {torch.version.cuda}")
        print(f"✓ GPU: {torch.cuda.get_device_name(0)}")
    else:
        print("⚠️  No GPU detected (will use CPU)")
except Exception as e:
    print(f"❌ Error: {e}")

# Test 3: Critical imports
print("\n[3/7] Critical Dependencies")
deps = {
    'gradio': 'Gradio',
    'numpy': 'NumPy',
    'scipy': 'SciPy',
    'matplotlib': 'Matplotlib',
    'trimesh': 'Trimesh',
    'einops': 'Einops',
    'clip': 'OpenAI CLIP'
}

for module, name in deps.items():
    try:
        __import__(module)
        print(f"✓ {name}")
    except ImportError as e:
        print(f"❌ {name} - NOT INSTALLED")

# Test 4: Model checkpoints
print("\n[4/7] Model Checkpoints")
checkpoints_dir = './checkpoints'
dataset_name = 't2m'

if os.path.exists(checkpoints_dir):
    print(f"✓ Checkpoints directory exists: {checkpoints_dir}")

    # Check for specific model directories
    models_to_check = [
        f'{checkpoints_dir}/{dataset_name}/t2m_nlayer8_nhead6_ld384_ff1024_cdp0.1_rvq6ns',
        f'{checkpoints_dir}/{dataset_name}/rvq_nq6_dc512_nc512_noshare_qdp0.2',
        f'{checkpoints_dir}/{dataset_name}/length_estimator',
    ]

    for model_path in models_to_check:
        if os.path.exists(model_path):
            # Count files
            files = []
            for root, dirs, filenames in os.walk(model_path):
                files.extend(filenames)
            print(f"✓ {os.path.basename(model_path)} ({len(files)} files)")
        else:
            print(f"❌ {os.path.basename(model_path)} - NOT FOUND")
else:
    print(f"❌ Checkpoints directory NOT FOUND: {checkpoints_dir}")
    print("   Models must be present for the app to work!")

# Test 5: Try loading app.py
print("\n[5/7] App.py Syntax")
try:
    with open('app.py', 'r', encoding='utf-8') as f:
        compile(f.read(), 'app.py', 'exec')
    print("✓ app.py syntax is valid")
except FileNotFoundError:
    print("❌ app.py not found")
except SyntaxError as e:
    print(f"❌ Syntax error: {e}")

# Test 6: Required files
print("\n[6/7] Required Files")
required = ['app.py', 'requirements.txt', 'README.md']
for file in required:
    if os.path.exists(file):
        size = os.path.getsize(file)
        print(f"✓ {file} ({size} bytes)")
    else:
        print(f"❌ {file} - NOT FOUND")

# Test 7: Disk space for outputs
print("\n[7/7] Output Directory")
output_dir = './gradio_outputs'
try:
    os.makedirs(output_dir, exist_ok=True)
    print(f"✓ Output directory ready: {output_dir}")
except Exception as e:
    print(f"❌ Error creating output directory: {e}")

# Summary
print("\n" + "=" * 80)
print("SUMMARY")
print("=" * 80)

# Check if ready
if os.path.exists(checkpoints_dir) and os.path.exists('app.py'):
    print("\n✅ Basic setup looks good!")
    print("\nNext steps:")
    print("1. Test locally: python app.py")
    print("2. Visit http://localhost:7860 in browser")
    print("3. Try a prompt and check for errors")
    print("4. If it works locally, redeploy to HF")
else:
    print("\n⚠️  Setup incomplete!")
    if not os.path.exists(checkpoints_dir):
        print("\n❌ Missing: Model checkpoints")
        print("   • Download models to ./checkpoints/")
        print("   • Or configure model download in app.py")

print("\n" + "=" * 80)