"""
Smoke Tests for Cosmos Transfer2.5
Validates model loading and world-to-world translation
"""
import sys
import os
import tempfile
from pathlib import Path

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


def test_transfer_model_loading():
    """Test that Transfer2.5 model can be loaded"""
    print("\n" + "=" * 60)
    print("TEST: Transfer2.5 Model Loading")
    print("=" * 60)

    import torch
    from cosmos.loaders import load_transfer_pipeline, get_device_info, clear_model_cache

    # Clear any cached models first
    clear_model_cache()

    device_info = get_device_info()
    print(f"Device: {device_info['name']}")
    print(f"VRAM: {device_info['free_vram_gb']:.2f} / {device_info['total_vram_gb']:.2f} GB")

    if device_info['free_vram_gb'] < 60:
        print("WARNING: Less than 60GB VRAM available")
        print("Transfer2.5 requires ~65GB VRAM and may fail to load")

    try:
        pipe = load_transfer_pipeline()
        print("Model loaded successfully!")
        print(f"Pipeline type: {type(pipe).__name__}")
        return True
    except Exception as e:
        print(f"FAILED: {e}")
        return False


def test_transfer_control_extraction():
    """Test control signal extraction methods"""
    print("\n" + "=" * 60)
    print("TEST: Transfer2.5 Control Extraction")
    print("=" * 60)

    from PIL import Image
    import numpy as np
    from cosmos.utils_video import extract_edges, extract_depth_map

    # Create test image
    test_img = Image.new("RGB", (320, 240), color=(128, 128, 128))

    try:
        # Test edge extraction
        edges = extract_edges(test_img)
        assert edges is not None, "Edge extraction returned None"
        assert edges.size == test_img.size, f"Edge size mismatch: {edges.size}"
        print(f"Edge extraction: OK (size: {edges.size})")

        # Test depth extraction
        depth = extract_depth_map(test_img)
        assert depth is not None, "Depth extraction returned None"
        assert depth.size == test_img.size, f"Depth size mismatch: {depth.size}"
        print(f"Depth extraction: OK (size: {depth.size})")

        print("TEST PASSED!")
        return True

    except Exception as e:
        print(f"FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_transfer_video_inference():
    """Test Transfer2.5 video-to-video inference"""
    print("\n" + "=" * 60)
    print("TEST: Transfer2.5 Video Inference")
    print("=" * 60)

    import torch
    from cosmos.utils_video import create_test_video, load_video_frames
    from cosmos.infer_transfer import transfer_video

    # Create test input video
    test_video = create_test_video(num_frames=17, width=320, height=240)
    output_path = tempfile.mktemp(suffix=".mp4")

    try:
        result = transfer_video(
            input_video=test_video,
            prompt="Transform to nighttime scene with city lights",
            control_type="blur",
            num_inference_steps=10,  # Minimum for speed
            guidance_scale=7.0,
            seed=42,
            output_path=output_path
        )

        # Validate output
        assert result is not None, "Result is None"
        assert os.path.exists(result['video_path']), "Output video does not exist"
        assert result['output_frames'] >= 8, f"Too few frames: {result['output_frames']}"

        # Check video file
        file_size = os.path.getsize(result['video_path'])
        assert file_size > 1000, f"Video file too small: {file_size} bytes"

        print(f"Output video: {result['video_path']}")
        print(f"Input frames: {result['input_frames']}")
        print(f"Output frames: {result['output_frames']}")
        print(f"Control type: {result['control_type']}")
        print(f"Inference time: {result['inference_time_s']}s")
        print("TEST PASSED!")

        return True

    except Exception as e:
        print(f"FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False

    finally:
        # Cleanup
        if os.path.exists(test_video):
            os.remove(test_video)


def test_transfer_style_consistency():
    """Test that style transfer maintains structural consistency"""
    print("\n" + "=" * 60)
    print("TEST: Transfer2.5 Style Consistency")
    print("=" * 60)

    from PIL import Image
    from cosmos.utils_video import compute_ssim, extract_edges

    # This test validates the SSIM utility for later consistency checks
    # Actual model inference consistency is tested in smoke_transfer_video_inference

    try:
        # Create two similar images
        img1 = Image.new("RGB", (100, 100), color=(128, 128, 128))
        img2 = Image.new("RGB", (100, 100), color=(130, 130, 130))  # Slightly different

        ssim = compute_ssim(img1, img2)
        print(f"SSIM between similar images: {ssim:.4f}")

        assert 0 <= ssim <= 1, f"SSIM out of range: {ssim}"
        assert ssim > 0.9, f"Similar images should have high SSIM: {ssim}"

        # Create very different image
        img3 = Image.new("RGB", (100, 100), color=(255, 0, 0))  # Red
        ssim_diff = compute_ssim(img1, img3)
        print(f"SSIM between different images: {ssim_diff:.4f}")

        assert ssim_diff < ssim, "Different images should have lower SSIM"

        print("TEST PASSED!")
        return True

    except Exception as e:
        print(f"FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False


def run_all_transfer_tests():
    """Run all Transfer2.5 smoke tests"""
    print("\n" + "=" * 60)
    print("COSMOS TRANSFER2.5 SMOKE TESTS")
    print("=" * 60)

    results = {}

    # Test 1: Control extraction (no model needed)
    results['control_extraction'] = test_transfer_control_extraction()

    # Test 2: Style consistency utilities (no model needed)
    results['style_consistency'] = test_transfer_style_consistency()

    # Test 3: Model loading (requires ~65GB VRAM)
    results['model_loading'] = test_transfer_model_loading()

    # Test 4: Video inference (requires loaded model)
    if results['model_loading']:
        results['video_inference'] = test_transfer_video_inference()
    else:
        results['video_inference'] = False
        print("\nSkipping inference test due to model loading failure")

    # Summary
    print("\n" + "=" * 60)
    print("TRANSFER2.5 TEST SUMMARY")
    print("=" * 60)

    passed = sum(1 for v in results.values() if v)
    total = len(results)

    for test_name, passed_flag in results.items():
        status = "PASSED" if passed_flag else "FAILED"
        print(f"  {test_name}: {status}")

    print(f"\nTotal: {passed}/{total} tests passed")

    return passed == total


if __name__ == "__main__":
    success = run_all_transfer_tests()
    sys.exit(0 if success else 1)