"""
Debug test for Cosmos Pipeline - isolate the pipeline bottleneck
"""

import torch
import time
import pytest

@pytest.fixture
def pipeline():
    return test_pipeline_init()

def test_pipeline_init():
    """Test only pipeline construction"""
    print("Testing Pipeline init...")
    
    from mimic_video.cosmos_predict import (
        TINY_VAE_CONFIG, TINY_TRANSFORMER_CONFIG, TINY_T5_CONFIG
    )
    from diffusers.models.autoencoders.autoencoder_kl_cosmos import AutoencoderKLCosmos
    from diffusers.models.transformers.transformer_cosmos import CosmosTransformer3DModel
    from diffusers.schedulers.scheduling_edm_euler import EDMEulerScheduler
    from diffusers import CosmosVideoToWorldPipeline
    from transformers import T5EncoderModel, T5TokenizerFast, T5Config
    
    start = time.time()
    transformer = CosmosTransformer3DModel(num_layers=2, **TINY_TRANSFORMER_CONFIG)
    print(f"  Transformer init: {time.time() - start:.2f}s")
    
    start = time.time()
    vae = AutoencoderKLCosmos(**TINY_VAE_CONFIG)
    print(f"  VAE init: {time.time() - start:.2f}s")
    
    start = time.time()
    scheduler = EDMEulerScheduler()
    print(f"  Scheduler init: {time.time() - start:.2f}s")
    
    start = time.time()
    t5_config = T5Config(**TINY_T5_CONFIG)
    text_encoder = T5EncoderModel(t5_config)
    tokenizer = T5TokenizerFast.from_pretrained("google-t5/t5-small")
    print(f"  T5 init: {time.time() - start:.2f}s")
    
    # Dummy safety checker
    class DummySafetyChecker:
        def __call__(self, *args, **kwargs):
            return [], False
        def to(self, *args, **kwargs):
            return self
        def check_text_safety(self, text):
            return True
        def check_video_safety(self, video):
            return True
    
    start = time.time()
    pipeline = CosmosVideoToWorldPipeline(
        transformer=transformer,
        vae=vae,
        scheduler=scheduler,
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        safety_checker=DummySafetyChecker(),
    )
    print(f"  Pipeline construction: {time.time() - start:.2f}s")
    
    print("✓ Pipeline init passed!")
    return pipeline

def test_pipeline_call(pipeline):
    """Test the pipeline __call__ method"""
    print("\nTesting Pipeline __call__...")
    
    # Create a tiny video input
    video = torch.rand(1, 5, 3, 32, 32)  # B, T, C, H, W
    print(f"Input video shape: {video.shape}")
    
    start = time.time()
    with torch.inference_mode():
        result = pipeline(
            video=video,
            prompt="test",
            num_inference_steps=1,
            guidance_scale=1.0,
            output_type="latent",
        )
    print(f"Pipeline call took: {time.time() - start:.2f}s")
    
    print("✓ Pipeline call passed!")
    return result

if __name__ == "__main__":
    pipeline = test_pipeline_init()
    test_pipeline_call(pipeline)
    print("\n✓ All pipeline tests passed!")
