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

from loops.prompt_nanobanana import PromptNanobanana


@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 image prompts for Nanobanana.")
        prompt_path = tmp_file.name
        
    yield prompt_path
    
    if os.path.exists(prompt_path):
        os.unlink(prompt_path)


@pytest.fixture
def sample_concept():
    """Return a sample ad concept string."""
    return "A vibrant ad featuring a modern smartphone with sleek design"


@patch('loops.prompt_nanobanana.ChatOpenAI')
def test_prompt_nanobanana_initialization(mock_openai, fake_system_prompt):
    """Test that PromptNanobanana initializes correctly."""
    # Mock the ChatOpenAI client
    mock_client = Mock()
    mock_openai.return_value = mock_client
    
    with patch.object(PromptNanobanana, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        prompt_gen = PromptNanobanana()
        
        assert prompt_gen.model == mock_client
        mock_openai.assert_called_once_with(model="gpt-4.1")


@patch('loops.prompt_nanobanana.ChatOpenAI')
def test_generate_returns_string(mock_openai, sample_concept, fake_system_prompt):
    """Test that generate returns a string."""
    # Mock the ChatOpenAI client and its response
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Generated image prompt for Nanobanana"
    mock_client.invoke.return_value = mock_response
    mock_openai.return_value = mock_client
    
    with patch.object(PromptNanobanana, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        prompt_gen = PromptNanobanana()
        result = prompt_gen.generate_from(sample_concept)
        
        assert isinstance(result, str)
        assert result == "Generated image prompt for Nanobanana"


@patch('loops.prompt_nanobanana.ChatOpenAI')
def test_generate_calls_model_correctly(mock_openai, sample_concept, fake_system_prompt):
    """Test that generate 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 image prompt"
    mock_client.invoke.return_value = mock_response
    mock_openai.return_value = mock_client
    
    with patch.object(PromptNanobanana, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        prompt_gen = PromptNanobanana()
        result = prompt_gen.generate_from(sample_concept)
        
        # 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 image prompt"


@patch('loops.prompt_nanobanana.ChatOpenAI')
def test_construct_full_prompt_structure(mock_openai, sample_concept, fake_system_prompt):
    """Test that _construct_full_prompt returns correct structure."""
    mock_client = Mock()
    mock_openai.return_value = mock_client
    
    with patch.object(PromptNanobanana, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        prompt_gen = PromptNanobanana()
        content = prompt_gen._construct_full_prompt(sample_concept)
        
        # Should return a list with 2 items: system_prompt and concept
        assert isinstance(content, list)
        assert len(content) == 2
        
        # First item is system prompt (dict with type and text)
        assert isinstance(content[0], dict)
        assert content[0]["type"] == "text"
        
        # Second item is the concept (dict with type and text)
        assert isinstance(content[1], dict)
        assert content[1]["type"] == "text"
        assert content[1]["text"] == sample_concept


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


@patch('loops.prompt_nanobanana.ChatOpenAI')
def test_empty_concept_handling(mock_openai, fake_system_prompt):
    """Test that PromptNanobanana handles empty concept string."""
    mock_client = Mock()
    mock_openai.return_value = mock_client
    
    with patch.object(PromptNanobanana, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        prompt_gen = PromptNanobanana()
        content = prompt_gen._construct_full_prompt("")
        
        # Should still create proper structure with empty text
        assert len(content) == 2
        assert content[1]["text"] == ""


@patch('loops.prompt_nanobanana.ChatOpenAI')
def test_long_concept_handling(mock_openai, fake_system_prompt):
    """Test that PromptNanobanana handles long concept strings."""
    long_concept = "A " * 1000  # Very long string
    
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Generated prompt for long concept"
    mock_client.invoke.return_value = mock_response
    mock_openai.return_value = mock_client
    
    with patch.object(PromptNanobanana, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        prompt_gen = PromptNanobanana()
        result = prompt_gen.generate_from(long_concept)
        
        # Should successfully generate without errors
        assert result == "Generated prompt for long concept"


@patch('loops.prompt_nanobanana.ChatOpenAI')
def test_concept_with_special_characters(mock_openai, fake_system_prompt):
    """Test that PromptNanobanana handles concepts with special characters."""
    special_concept = "Ad with émojis 🎉 and symbols @#$%^&*()"
    
    mock_client = Mock()
    mock_openai.return_value = mock_client
    
    with patch.object(PromptNanobanana, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        prompt_gen = PromptNanobanana()
        content = prompt_gen._construct_full_prompt(special_concept)
        
        assert content[1]["text"] == special_concept


@patch('loops.prompt_nanobanana.ChatOpenAI')
def test_multiple_generate_calls(mock_openai, sample_concept, fake_system_prompt):
    """Test that generate can be called multiple times."""
    mock_client = Mock()
    mock_response = Mock()
    mock_response.content = "Test prompt"
    mock_client.invoke.return_value = mock_response
    mock_openai.return_value = mock_client
    
    with patch.object(PromptNanobanana, 'SYSTEM_PROMPT_PATH', fake_system_prompt):
        prompt_gen = PromptNanobanana()
        
        result1 = prompt_gen.generate_from(sample_concept)
        result2 = prompt_gen.generate_from("Different concept")
        
        # Both calls should succeed
        assert result1 == "Test prompt"
        assert result2 == "Test prompt"
        
        # Model should be invoked twice
        assert mock_client.invoke.call_count == 2
