import pytest
import tempfile
import os
from unittest.mock import Mock, patch
from PIL import Image

from loops.new_ad_visual import NewAdVisual


@pytest.fixture
def fake_product_image():
    """Generate a temporary product image file for testing."""
    with tempfile.NamedTemporaryFile(mode='wb', suffix='.jpg', delete=False) as tmp_file:
        jpeg_path = tmp_file.name
        img = Image.new('RGB', (200, 200), color='green')
        img.save(jpeg_path, 'JPEG')
        
    yield jpeg_path
    
    if os.path.exists(jpeg_path):
        os.unlink(jpeg_path)


@pytest.fixture
def sample_prompt():
    """Return a sample image generation prompt."""
    return "Generate a vibrant advertisement featuring the product on a beach at sunset"


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_new_ad_visual_initialization(mock_gemini, fake_product_image):
    """Test that NewAdVisual initializes correctly."""
    # Mock the ChatGoogleGenerativeAI client
    mock_client = Mock()
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    
    assert visual_generator.product_image_path == fake_product_image
    assert visual_generator.model == mock_client
    mock_gemini.assert_called_once_with(model="gemini-2.5-flash-image")


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_generate_following_returns_string(mock_gemini, fake_product_image, sample_prompt):
    """Test that generate_following returns a string."""
    # Mock the ChatGoogleGenerativeAI client and its response
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Generated visual description or URL"
    mock_client.invoke.return_value = mock_response
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    result = visual_generator.generate_following(sample_prompt)
    
    assert isinstance(result, str)
    assert result == "Generated visual description or URL"


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_generate_following_calls_model_correctly(mock_gemini, fake_product_image, sample_prompt):
    """Test that generate_following calls the model with correct structure."""
    from langchain_core.messages import HumanMessage
    
    # Mock the ChatGoogleGenerativeAI client
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Test visual"
    mock_client.invoke.return_value = mock_response
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    result = visual_generator.generate_following(sample_prompt)
    
    # Verify model.invoke was called
    assert mock_client.invoke.called
    call_args = mock_client.invoke.call_args[0][0]
    
    # Check that it was called with a list containing a HumanMessage
    assert isinstance(call_args, list)
    assert len(call_args) == 1
    assert isinstance(call_args[0], HumanMessage)
    
    # Verify result is the model output
    assert result == "Test visual"


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_construct_full_prompt_structure(mock_gemini, fake_product_image, sample_prompt):
    """Test that _construct_full_prompt returns correct structure."""
    mock_client = Mock()
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    content = visual_generator._construct_full_prompt(sample_prompt)
    
    # Should return a list with 2 items: product_image and prompt
    assert isinstance(content, list)
    assert len(content) == 2
    
    # First item is product image (dict with type and image_url)
    assert isinstance(content[0], dict)
    assert content[0]["type"] == "image_url"
    assert "image_url" in content[0]
    assert "url" in content[0]["image_url"]
    
    # Second item is the prompt text (dict with type and text)
    assert isinstance(content[1], dict)
    assert content[1]["type"] == "text"
    assert content[1]["text"] == sample_prompt


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_model_attribute_is_correct_type(mock_gemini, fake_product_image):
    """Test that the model attribute is set correctly."""
    mock_client = Mock()
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    
    # Verify the model is the mocked ChatGoogleGenerativeAI instance
    assert visual_generator.model == mock_client


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_new_ad_visual_with_nonexistent_image(mock_gemini, sample_prompt):
    """Test that NewAdVisual handles non-existent product image during generation."""
    mock_client = Mock()
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual("/nonexistent/product.jpg")
    
    # Should raise FileNotFoundError when trying to generate
    with pytest.raises(FileNotFoundError):
        visual_generator.generate_following(sample_prompt)


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_empty_prompt_handling(mock_gemini, fake_product_image):
    """Test that NewAdVisual handles empty prompt string."""
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Generated with empty prompt"
    mock_client.invoke.return_value = mock_response
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    result = visual_generator.generate_following("")
    
    # Should still work with empty prompt
    assert result == "Generated with empty prompt"
    content = visual_generator._construct_full_prompt("")
    assert content[1]["text"] == ""


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_long_prompt_handling(mock_gemini, fake_product_image):
    """Test that NewAdVisual handles very long prompts."""
    long_prompt = "Create an image " * 500  # Very long string
    
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Generated with long prompt"
    mock_client.invoke.return_value = mock_response
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    result = visual_generator.generate_following(long_prompt)
    
    # Should successfully generate without errors
    assert result == "Generated with long prompt"
    content = visual_generator._construct_full_prompt(long_prompt)
    assert content[1]["text"] == long_prompt


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_prompt_with_special_characters(mock_gemini, fake_product_image):
    """Test that NewAdVisual handles prompts with special characters."""
    special_prompt = "Create ad with émojis 🎉🚀 and symbols @#$%^&*()"
    
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Generated with special chars"
    mock_client.invoke.return_value = mock_response
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    result = visual_generator.generate_following(special_prompt)
    
    assert result == "Generated with special chars"
    content = visual_generator._construct_full_prompt(special_prompt)
    assert content[1]["text"] == special_prompt


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_multiple_generate_calls(mock_gemini, fake_product_image, sample_prompt):
    """Test that generate_following can be called multiple times."""
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Test visual"
    mock_client.invoke.return_value = mock_response
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    
    result1 = visual_generator.generate_following(sample_prompt)
    result2 = visual_generator.generate_following("Different prompt")
    
    # Both calls should succeed
    assert result1 == "Test visual"
    assert result2 == "Test visual"
    
    # Model should be invoked twice
    assert mock_client.invoke.call_count == 2


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_different_image_formats(mock_gemini):
    """Test that NewAdVisual works with different image file extensions."""
    # Create PNG image
    with tempfile.NamedTemporaryFile(mode='wb', suffix='.png', delete=False) as tmp_file:
        png_path = tmp_file.name
        img = Image.new('RGB', (100, 100), color='red')
        img.save(png_path, 'PNG')
    
    try:
        mock_client = Mock()
        mock_response = Mock()
        mock_response.content = "Generated from PNG"
        mock_client.invoke.return_value = mock_response
        mock_gemini.return_value = mock_client
        
        visual_generator = NewAdVisual(png_path)
        result = visual_generator.generate_following("Test prompt")
        
        assert result == "Generated from PNG"
        assert visual_generator.product_image_path == png_path
    finally:
        if os.path.exists(png_path):
            os.unlink(png_path)


@patch('loops.new_ad_visual.ChatGoogleGenerativeAI')
def test_product_image_path_stored_correctly(mock_gemini, fake_product_image):
    """Test that product_image_path attribute stores the path correctly."""
    mock_client = Mock()
    mock_gemini.return_value = mock_client
    
    visual_generator = NewAdVisual(fake_product_image)
    
    assert visual_generator.product_image_path == fake_product_image
    assert isinstance(visual_generator.product_image_path, str)
