#!/usr/bin/env python3
"""
Test runner for FBMC notebooks on HuggingFace Space
Executes notebooks and reports results
"""
import subprocess
import sys
import os
from pathlib import Path
from datetime import datetime

def run_notebook(notebook_path, timeout=600):
    """Execute a Jupyter notebook and return success status"""
    print(f"\n{'='*60}")
    print(f"Running: {notebook_path.name}")
    print(f"Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'='*60}\n")

    cmd = [
        "jupyter", "nbconvert",
        "--to", "notebook",
        "--execute",
        "--inplace",
        f"--ExecutePreprocessor.timeout={timeout}",
        str(notebook_path)
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout+60)

        if result.returncode == 0:
            print(f"\n✅ SUCCESS: {notebook_path.name}")
            return True
        else:
            print(f"\n❌ FAILED: {notebook_path.name}")
            print(f"Error: {result.stderr}")
            return False
    except subprocess.TimeoutExpired:
        print(f"\n❌ TIMEOUT: {notebook_path.name} exceeded {timeout}s")
        return False
    except Exception as e:
        print(f"\n❌ ERROR: {str(e)}")
        return False

def main():
    # Check environment
    if "HF_TOKEN" not in os.environ:
        print("⚠️  Warning: HF_TOKEN not set!")
        print("   Set with: export HF_TOKEN='your_token'\n")

    # Check GPU
    try:
        import torch
        print(f"GPU Available: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"GPU Name: {torch.cuda.get_device_name(0)}")
            print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    except ImportError:
        print("⚠️  PyTorch not installed")

    print("\n")

    # Define test sequence
    tests = [
        ("/data/inference_smoke_test.ipynb", 300),  # 5 min
        # Uncomment to run full tests:
        # ("/data/inference_full_14day.ipynb", 900),   # 15 min
        # ("/data/evaluation.ipynb", 300),             # 5 min
    ]

    results = {}

    for notebook_path, timeout in tests:
        nb_path = Path(notebook_path)
        if not nb_path.exists():
            print(f"❌ Not found: {notebook_path}")
            results[nb_path.name] = False
            continue

        success = run_notebook(nb_path, timeout)
        results[nb_path.name] = success

    # Summary
    print(f"\n{'='*60}")
    print("SUMMARY")
    print(f"{'='*60}")
    for name, success in results.items():
        status = "✅ PASS" if success else "❌ FAIL"
        print(f"{status}: {name}")

    total = len(results)
    passed = sum(results.values())
    print(f"\nResults: {passed}/{total} passed")

    return 0 if all(results.values()) else 1

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