#!/usr/bin/env python3
"""
Repro for multimodal vision batch inference bug reported in #56125.

Environment:
ray                                2.49.1
vllm                               0.10.1.1

Usage:
python reproduce_multimodal_bug.py

Reproduces
1. WITH system prompt - should work (image_sizes present, reasonable responses)
2. WITHOUT system prompt - should fail (image_sizes missing, nonsensical responses)

"""

import ray
from ray.data.llm import build_llm_processor, vLLMEngineProcessorConfig


def test_with_system_prompt():
    """Test case WITH system prompt - should work correctly."""
    print("=" * 60)
    print("TESTING WITH SYSTEM PROMPT (should work)")
    print("=" * 60)
    
    config = vLLMEngineProcessorConfig(
        model_source="HuggingFaceTB/SmolVLM-256M-Instruct",
        task_type="generate",
        engine_kwargs=dict(
            # Skip CUDA graph capturing to reduce startup time.
            enforce_eager=True,
            dtype="half",
        ),
        apply_chat_template=True,
        has_image=True,
        tokenize=True,
        detokenize=True,
        batch_size=16,
        accelerator_type="L4",
        concurrency=1,
    )

    processor = build_llm_processor(
        config,
        preprocess=lambda row: dict(
            messages=[
                {"role": "system", "content": "You are an assistant"},  # STRING content
                {
                    "role": "user",
                    "content": [  # LIST content - creates mixed content types
                        {
                            "type": "image",
                            "image": "https://i.scdn.co/image/ab67616d0000b27331efa547b5c271f00bcde9b9",
                        },
                        {
                            "type": "text",
                            "text": f"Can you describe this image in {row['id']} words?",
                        },
                    ],
                },
            ],
            sampling_params=dict(
                temperature=0.3,
                max_tokens=50,
            ),
        ),
        postprocess=lambda row: {
            "resp": row["generated_text"],
            **row,
        },
    )

    ds = ray.data.range(3)
    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
    ds = processor(ds)
    ds = ds.materialize()
    outs = ds.take_all()
    
    print(f"Results: {len(outs)} items")
    for i, out in enumerate(outs):
        print(f"\nResult {i}:")
        print(f"  image_sizes: {out.get('image_sizes', 'MISSING')}")
        print(f"  messages preserved: {len(out.get('messages', []))} messages")
        if 'messages' in out and len(out['messages']) > 1:
            user_content = out['messages'][1]['content']
            print(f"  user content type: {type(user_content)}")
            if hasattr(user_content, '__len__') and len(user_content) > 0:
                first_item = user_content[0]
                print(f"  first content item: {first_item}")
        print(f"  response: {out['resp'][:100]}...")
    
    return outs


def test_without_system_prompt():
    """Test case WITHOUT system prompt - should fail (before fix)."""
    print("\n" + "=" * 60)
    print("TESTING WITHOUT SYSTEM PROMPT (should fail before fix)")
    print("=" * 60)
    
    config = vLLMEngineProcessorConfig(
        model_source="HuggingFaceTB/SmolVLM-256M-Instruct",
        task_type="generate",
        engine_kwargs=dict(
            # Skip CUDA graph capturing to reduce startup time.
            enforce_eager=True,
            dtype="half",
        ),
        apply_chat_template=True,
        has_image=True,
        tokenize=True,
        detokenize=True,
        batch_size=16,
        accelerator_type="L4",
        concurrency=1,
    )

    processor = build_llm_processor(
        config,
        preprocess=lambda row: dict(
            messages=[
                # NO system prompt - only user message with LIST content
                {
                    "role": "user",
                    "content": [  # Only LIST content - uniform content types
                        {
                            "type": "image",
                            "image": "https://i.scdn.co/image/ab67616d0000b27331efa547b5c271f00bcde9b9",
                        },
                        {
                            "type": "text",
                            "text": f"Can you describe this image in {row['id']} words?",
                        },
                    ],
                },
            ],
            sampling_params=dict(
                temperature=0.3,
                max_tokens=50,
            ),
        ),
        postprocess=lambda row: {
            "resp": row["generated_text"],
            **row,
        },
    )

    ds = ray.data.range(3)
    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
    ds = processor(ds)
    ds = ds.materialize()
    outs = ds.take_all()
    
    print(f"Results: {len(outs)} items")
    for i, out in enumerate(outs):
        print(f"\nResult {i}:")
        print(f"  image_sizes: {out.get('image_sizes', 'MISSING')}")
        print(f"  messages preserved: {len(out.get('messages', []))} messages")
        if 'messages' in out and len(out['messages']) > 0:
            user_content = out['messages'][0]['content']
            print(f"  user content type: {type(user_content)}")
            if hasattr(user_content, '__len__') and len(user_content) > 0:
                first_item = user_content[0]
                print(f"  first content item: {first_item}")
        print(f"  response: {out['resp'][:100]}...")
    
    return outs


def analyze_results(with_system, without_system):
    """Analyze and compare the results."""
    print("\n" + "=" * 60)
    print("ANALYSIS")
    print("=" * 60)
    
    with_result = with_system[0]
    without_result = without_system[0]
    
    print(f"\nWITH system prompt:")
    print(f"   image_sizes present: {'image_sizes' in with_result}")
    print(f"   image_sizes value: {with_result.get('image_sizes', 'MISSING')}")
    print(f"   response preview: {with_result['resp'][:80]}...")
    
    print(f"\nWITHOUT system prompt:")
    print(f"   image_sizes present: {'image_sizes' in without_result}")
    print(f"   image_sizes value: {without_result.get('image_sizes', 'MISSING')}")
    print(f"   response preview: {without_result['resp'][:80]}...")
    
    # Check if bug is reproduced
    bug_reproduced = (
        'image_sizes' in with_result and 
        'image_sizes' not in without_result
    )
    
    print(f"\nBug reproduction: {'CONFIRMED' if bug_reproduced else 'NOT REPRODUCED'}")
    
    if bug_reproduced:
        print("   The bug is reproduced exactly as reported:")
        print("   - WITH system prompt: images work, image_sizes present")
        print("   - WITHOUT system prompt: images fail, image_sizes missing")
        print("   - This confirms our PyArrow vs Python object theory")
    else:
        print("   The bug was not reproduced - either:")
        print("   - The fix is already applied and working")
        print("   - Our theory needs refinement")
    
    return bug_reproduced


def main():
    """Run the full reproduction test."""
    print("REPRODUCING MULTIMODAL VISION BATCH INFERENCE BUG")
    print("Issue #56125: [Data] multimodal vision batch inference skips the image preparation step")
    
    try:
        # Initialize Ray
        ray.init(ignore_reinit_error=True)
        
        # Run both test cases
        print("\nRunning test cases...")
        with_system_results = test_with_system_prompt()
        without_system_results = test_without_system_prompt()
        
        # Analyze results
        bug_reproduced = analyze_results(with_system_results, without_system_results)
        
        print(f"\nCONCLUSION:")
        if bug_reproduced:
            print("   The bug has been reproduced!")
            print("   Can now apply the fix and run this script again to verify it works.")
        else:
            print("   Bug was not reproduced - the fix may be already applied.")
            
    except Exception as e:
        print(f"Error during reproduction: {e}")
        import traceback
        traceback.print_exc()
    finally:
        ray.shutdown()


if __name__ == "__main__":
    main()