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

from loops.new_ad_concept import NewAdConcept


@pytest.fixture
def fake_brand_brief_pdf():
    """Generate a temporary PDF file for testing."""
    from reportlab.pdfgen import canvas
    from reportlab.lib.pagesizes import letter
    
    with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp_file:
        pdf_path = tmp_file.name
        c = canvas.Canvas(pdf_path, pagesize=letter)
        c.drawString(100, 750, "Brand Brief: Test Product")
        c.save()
        
    yield pdf_path
    
    if os.path.exists(pdf_path):
        os.unlink(pdf_path)


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


@pytest.fixture
def fake_system_prompt():
    """Create a temporary system prompt file."""
    with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as tmp_file:
        tmp_file.write("You are an AI that generates ad concepts.")
        prompt_path = tmp_file.name
        
    yield prompt_path
    
    if os.path.exists(prompt_path):
        os.unlink(prompt_path)


@patch('loops.new_ad_concept.ChatOpenAI')
def test_new_ad_concept_initialization(mock_openai, fake_brand_brief_pdf, fake_system_prompt):
    """Test that NewAdConcept initializes correctly."""
    # Mock the ChatOpenAI client
    mock_client = Mock()
    mock_openai.return_value = mock_client
    
    # Temporarily replace the system prompt path
    with patch.object(NewAdConcept, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        concept_generator = NewAdConcept(fake_brand_brief_pdf)
        
        assert concept_generator.brand_brief_path == fake_brand_brief_pdf
        assert concept_generator.model == mock_client
        mock_openai.assert_called_once_with(model="o3")


@patch('loops.new_ad_concept.ChatOpenAI')
def test_generate_from_returns_string(
    mock_openai, 
    fake_brand_brief_pdf, 
    fake_inspiration_jpeg, 
    fake_system_prompt
):
    """Test that generate_from returns a string."""
    # Mock the ChatOpenAI client and its response
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Generated ad concept"
    mock_client.invoke.return_value = mock_response
    mock_openai.return_value = mock_client
    
    with patch.object(NewAdConcept, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        concept_generator = NewAdConcept(fake_brand_brief_pdf)
        result = concept_generator.generate_from(fake_inspiration_jpeg)
        
        assert isinstance(result, str)
        assert result == "Generated ad concept"


@patch('loops.new_ad_concept.ChatOpenAI')
def test_generate_from_calls_model_correctly(
    mock_openai, 
    fake_brand_brief_pdf, 
    fake_inspiration_jpeg,
    fake_system_prompt
):
    """Test that generate_from calls the model with correct structure."""
    from langchain_core.messages import HumanMessage
    
    # Mock the ChatOpenAI client
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Test concept"
    mock_client.invoke.return_value = mock_response
    mock_openai.return_value = mock_client
    
    with patch.object(NewAdConcept, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        concept_generator = NewAdConcept(fake_brand_brief_pdf)
        result = concept_generator.generate_from(fake_inspiration_jpeg)
        
        # 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 concept"


@patch('loops.new_ad_concept.ChatOpenAI')
def test_construct_full_prompt_structure(
    mock_openai, 
    fake_brand_brief_pdf, 
    fake_inspiration_jpeg,
    fake_system_prompt
):
    """Test that _construct_full_prompt returns correct structure."""
    mock_client = Mock()
    mock_openai.return_value = mock_client
    
    with patch.object(NewAdConcept, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        concept_generator = NewAdConcept(fake_brand_brief_pdf)
        content = concept_generator._construct_full_prompt(fake_inspiration_jpeg)
        
        # Should return a list with 3 items: system_prompt, brief, inspiration
        assert isinstance(content, list)
        assert len(content) == 3
        
        # First item is system prompt (dict with type and text)
        assert isinstance(content[0], dict)
        assert content[0]["type"] == "text"
        
        # Second item is brand brief (dict with type and text)
        assert isinstance(content[1], dict)
        assert content[1]["type"] == "text"
        
        # Third item is inspiration ad (dict with type and image_url)
        assert isinstance(content[2], dict)
        assert content[2]["type"] == "image_url"


@patch('loops.new_ad_concept.ChatOpenAI')
def test_new_ad_concept_with_nonexistent_brief(mock_openai, fake_system_prompt):
    """Test that NewAdConcept handles non-existent brief file gracefully during generate."""
    mock_client = Mock()
    mock_openai.return_value = mock_client
    
    with patch.object(NewAdConcept, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        concept_generator = NewAdConcept("/nonexistent/brief.pdf")
        
        # Should raise FileNotFoundError when trying to generate
        with pytest.raises(FileNotFoundError):
            concept_generator.generate_from("/some/inspiration.jpg")


@patch('loops.new_ad_concept.ChatOpenAI')
def test_model_attribute_is_correct_type(mock_openai, fake_brand_brief_pdf, fake_system_prompt):
    """Test that the model attribute is set correctly."""
    mock_client = Mock()
    mock_openai.return_value = mock_client
    
    with patch.object(NewAdConcept, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        concept_generator = NewAdConcept(fake_brand_brief_pdf)
        
        # Verify the model is the mocked ChatOpenAI instance
        assert concept_generator.model == mock_client


@patch('loops.new_ad_concept.ChatOpenAI')
def test_system_prompt_loaded(mock_openai, fake_brand_brief_pdf, fake_system_prompt):
    """Test that system prompt is loaded correctly."""
    mock_client = Mock()
    mock_openai.return_value = mock_client
    
    with patch.object(NewAdConcept, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        concept_generator = NewAdConcept(fake_brand_brief_pdf)
        
        # Verify system prompt is loaded and has correct structure
        assert concept_generator.system_prompt is not None
        assert isinstance(concept_generator.system_prompt, dict)
        assert concept_generator.system_prompt["type"] == "text"
        assert "AI that generates ad concepts" in concept_generator.system_prompt["text"]
