"""
Cosmos Smoke Tests - Master Runner
Runs all smoke tests for Predict2.5 and Transfer2.5
"""
import sys
import os
import time
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))


def print_header(text):
    print("\n" + "=" * 70)
    print(text)
    print("=" * 70)


def get_system_info():
    """Get system and GPU information"""
    import torch

    info = {
        "python_version": sys.version.split()[0],
        "torch_version": torch.__version__,
        "cuda_available": torch.cuda.is_available(),
        "cuda_version": torch.version.cuda if torch.cuda.is_available() else "N/A",
    }

    if torch.cuda.is_available():
        device = torch.cuda.current_device()
        props = torch.cuda.get_device_properties(device)
        free_mem, total_mem = torch.cuda.mem_get_info(device)

        info.update({
            "gpu_name": props.name,
            "gpu_total_vram_gb": round(total_mem / (1024**3), 2),
            "gpu_free_vram_gb": round(free_mem / (1024**3), 2),
            "compute_capability": f"{props.major}.{props.minor}"
        })

    return info


def run_all_tests(skip_transfer=False):
    """
    Run all smoke tests

    Args:
        skip_transfer: Skip Transfer2.5 tests if VRAM is limited
    """
    print_header("NVIDIA COSMOS WORLD FOUNDATION MODELS - SMOKE TESTS")

    start_time = time.time()

    # System info
    print("\n[System Information]")
    sys_info = get_system_info()
    for key, value in sys_info.items():
        print(f"  {key}: {value}")

    # Check VRAM
    if sys_info.get('gpu_free_vram_gb', 0) < 30:
        print("\nWARNING: Less than 30GB VRAM available.")
        print("Predict2.5 tests may fail.")

    if sys_info.get('gpu_free_vram_gb', 0) < 60:
        print("\nWARNING: Less than 60GB VRAM available.")
        print("Transfer2.5 tests will likely fail (requires ~65GB).")
        if not skip_transfer:
            print("Consider running with --skip-transfer flag.")

    results = {
        "predict": {},
        "transfer": {}
    }

    # Run Predict2.5 tests
    print_header("PART 1: COSMOS PREDICT2.5 TESTS")

    from tests.smoke_predict import (
        test_predict_output_validation,
        test_predict_model_loading,
        test_predict_text2world_inference
    )

    results['predict']['output_validation'] = test_predict_output_validation()
    results['predict']['model_loading'] = test_predict_model_loading()

    if results['predict']['model_loading']:
        results['predict']['text2world_inference'] = test_predict_text2world_inference()
    else:
        results['predict']['text2world_inference'] = False
        print("\nSkipping Predict2.5 inference test - model failed to load")

    # Clear Predict model before Transfer tests
    if not skip_transfer:
        print("\nClearing Predict model from memory...")
        from cosmos.loaders import clear_model_cache
        clear_model_cache()

        # Run Transfer2.5 tests
        print_header("PART 2: COSMOS TRANSFER2.5 TESTS")

        from tests.smoke_transfer import (
            test_transfer_control_extraction,
            test_transfer_style_consistency,
            test_transfer_model_loading,
            test_transfer_video_inference
        )

        results['transfer']['control_extraction'] = test_transfer_control_extraction()
        results['transfer']['style_consistency'] = test_transfer_style_consistency()
        results['transfer']['model_loading'] = test_transfer_model_loading()

        if results['transfer']['model_loading']:
            results['transfer']['video_inference'] = test_transfer_video_inference()
        else:
            results['transfer']['video_inference'] = False
            print("\nSkipping Transfer2.5 inference test - model failed to load")
    else:
        print_header("PART 2: COSMOS TRANSFER2.5 TESTS (SKIPPED)")
        print("Transfer2.5 tests skipped due to --skip-transfer flag")
        results['transfer'] = {"skipped": True}

    # Final Summary
    total_time = time.time() - start_time

    print_header("FINAL TEST SUMMARY")

    print("\n[Predict2.5 Results]")
    predict_passed = 0
    predict_total = 0
    for test_name, passed in results['predict'].items():
        status = "PASSED" if passed else "FAILED"
        print(f"  {test_name}: {status}")
        predict_total += 1
        if passed:
            predict_passed += 1

    print(f"\n  Predict2.5: {predict_passed}/{predict_total} tests passed")

    if not skip_transfer and 'skipped' not in results['transfer']:
        print("\n[Transfer2.5 Results]")
        transfer_passed = 0
        transfer_total = 0
        for test_name, passed in results['transfer'].items():
            status = "PASSED" if passed else "FAILED"
            print(f"  {test_name}: {status}")
            transfer_total += 1
            if passed:
                transfer_passed += 1

        print(f"\n  Transfer2.5: {transfer_passed}/{transfer_total} tests passed")
    else:
        transfer_passed = 0
        transfer_total = 0
        print("\n[Transfer2.5 Results]")
        print("  SKIPPED")

    # Overall
    total_passed = predict_passed + transfer_passed
    total_tests = predict_total + transfer_total

    print(f"\n[Overall]")
    print(f"  Total: {total_passed}/{total_tests} tests passed")
    print(f"  Time: {total_time:.2f}s")

    all_passed = (predict_passed == predict_total) and (
        skip_transfer or transfer_passed == transfer_total
    )

    if all_passed:
        print("\n  STATUS: ALL TESTS PASSED")
    else:
        print("\n  STATUS: SOME TESTS FAILED")

    return all_passed


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Run Cosmos smoke tests")
    parser.add_argument(
        "--skip-transfer",
        action="store_true",
        help="Skip Transfer2.5 tests (requires ~65GB VRAM)"
    )
    parser.add_argument(
        "--predict-only",
        action="store_true",
        help="Run only Predict2.5 tests"
    )

    args = parser.parse_args()

    skip_transfer = args.skip_transfer or args.predict_only

    success = run_all_tests(skip_transfer=skip_transfer)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()