"""
Tests for generator: NL → Ape code via OpenAI.
"""

import pytest
from unittest.mock import Mock, patch, MagicMock

from ape_openai.generator import (
    generate_ape_from_nl,
    validate_generated_ape,
    generate_and_compile_ape
)


@patch('ape_openai.generator.OPENAI_AVAILABLE', True)
@patch('ape_openai.generator.OpenAI')
def test_generate_ape_from_nl_success(mock_openai):
    """Test successful Ape code generation from natural language."""
    # Mock OpenAI API response
    mock_choice = Mock()
    mock_choice.message.content = """
```ape
task add:
    inputs: a: Integer, b: Integer
    outputs: sum: Integer
    constraints: a > 0, b > 0
    steps:
        sum = a + b
```
"""
    
    mock_response = Mock()
    mock_response.choices = [mock_choice]
    
    mock_client = Mock()
    mock_client.chat.completions.create.return_value = mock_response
    mock_openai.return_value = mock_client
    
    result = generate_ape_from_nl("Create a task that adds two positive numbers")
    
    assert "task add:" in result
    assert "inputs:" in result
    assert "outputs:" in result
    assert "constraints:" in result
    assert "steps:" in result
    assert "```" not in result  # Markdown fences should be removed


@patch('ape_openai.generator.OPENAI_AVAILABLE', True)
@patch('ape_openai.generator.OpenAI')
def test_generate_ape_from_nl_custom_model(mock_openai):
    """Test generation with custom model."""
    mock_choice = Mock()
    mock_choice.message.content = "task test:\n    inputs:\n    outputs:\n    steps:"
    
    mock_response = Mock()
    mock_response.choices = [mock_choice]
    
    mock_client = Mock()
    mock_client.chat.completions.create.return_value = mock_response
    mock_openai.return_value = mock_client
    
    generate_ape_from_nl("Create a test task", model="gpt-3.5-turbo")
    
    # Verify correct model was used
    call_args = mock_client.chat.completions.create.call_args
    assert call_args[1]["model"] == "gpt-3.5-turbo"


def test_validate_generated_ape_valid():
    """Test validation of valid Ape code."""
    valid_code = """
task add:
    inputs: a: Integer, b: Integer
    outputs: sum: Integer
    steps:
        sum = a + b
"""
    
    # Need to patch both imports inside the function
    with patch('ape.compile') as mock_compile:
        mock_compile.return_value = Mock()  # Successful compilation
        
        is_valid, error_msg = validate_generated_ape(valid_code)
        
        assert is_valid is True
        assert error_msg is None
        mock_compile.assert_called_once()


def test_validate_generated_ape_invalid():
    """Test validation of invalid Ape code."""
    invalid_code = "invalid ape code {"
    
    with patch('ape.compile') as mock_compile:
        mock_compile.side_effect = Exception("Syntax error at line 1")
        
        is_valid, error_msg = validate_generated_ape(invalid_code)
        
        assert is_valid is False
        assert error_msg is not None  # Just check there's an error message


@patch('ape_openai.generator.OPENAI_AVAILABLE', True)
@patch('ape_openai.generator.OpenAI')
@patch('ape.compile')
def test_generate_and_compile_ape_success(mock_compile, mock_openai):
    """Test successful generation and compilation."""
    # Mock OpenAI response
    mock_choice = Mock()
    mock_choice.message.content = """
task multiply:
    inputs: x: Integer, y: Integer
    outputs: result: Integer
    steps:
        result = x * y
"""
    
    mock_response = Mock()
    mock_response.choices = [mock_choice]
    
    mock_client = Mock()
    mock_client.chat.completions.create.return_value = mock_response
    mock_openai.return_value = mock_client
    
    # Mock successful compilation
    mock_module = Mock()
    mock_compile.return_value = mock_module
    
    code, module = generate_and_compile_ape("Multiply two integers")
    
    assert "task multiply:" in code
    assert module == mock_module


@patch('ape_openai.generator.OPENAI_AVAILABLE', True)
@patch('ape_openai.generator.OpenAI')
@patch('ape.compile')
def test_generate_and_compile_ape_retry(mock_compile, mock_openai):
    """Test generation with retry on compilation failure."""
    # Mock OpenAI responses (first invalid, second valid)
    invalid_choice = Mock()
    invalid_choice.message.content = "invalid code"
    
    valid_choice = Mock()
    valid_choice.message.content = "task test:\n    inputs:\n    outputs:\n    steps:"
    
    mock_response_1 = Mock()
    mock_response_1.choices = [invalid_choice]
    
    mock_response_2 = Mock()
    mock_response_2.choices = [valid_choice]
    
    mock_client = Mock()
    mock_client.chat.completions.create.side_effect = [mock_response_1, mock_response_2]
    mock_openai.return_value = mock_client
    
    # Mock compilation: 
    # 1st call - validation of "invalid code" fails
    # 2nd call - validation of "task test..." succeeds  
    # 3rd call - final compilation succeeds
    mock_module = Mock()
    mock_compile.side_effect = [
        Exception("Invalid syntax"),  # First validation fails
        mock_module,                   # Second validation succeeds
        mock_module                    # Final compilation
    ]
    
    code, module = generate_and_compile_ape("Create a test task", max_retries=2)
    
    assert code == "task test:\n    inputs:\n    outputs:\n    steps:"
    assert module == mock_module
    assert mock_client.chat.completions.create.call_count == 2


@patch('ape_openai.generator.OPENAI_AVAILABLE', True)
@patch('ape_openai.generator.OpenAI')
@patch('ape.compile')
def test_generate_and_compile_ape_max_retries_exceeded(mock_compile, mock_openai):
    """Test failure when max retries exceeded."""
    # Mock OpenAI response with invalid code
    mock_choice = Mock()
    mock_choice.message.content = "invalid code"
    
    mock_response = Mock()
    mock_response.choices = [mock_choice]
    
    mock_client = Mock()
    mock_client.chat.completions.create.return_value = mock_response
    mock_openai.return_value = mock_client
    
    # Mock compilation failure
    mock_compile.side_effect = Exception("Invalid syntax")
    
    with pytest.raises(Exception, match="Failed to generate valid Ape code after 2 attempts"):
        generate_and_compile_ape("Create invalid task", max_retries=2)


def test_validate_generated_ape_empty():
    """Test validation of empty code."""
    is_valid, error_msg = validate_generated_ape("")
    
    assert is_valid is False
    assert error_msg == "Empty code"


@patch('ape_openai.generator.OPENAI_AVAILABLE', True)
@patch('ape_openai.generator.OpenAI')
def test_generate_ape_from_nl_cleans_markdown(mock_openai):
    """Test that markdown code fences are properly removed."""
    mock_choice = Mock()
    mock_choice.message.content = """```ape
task test:
    inputs:
    outputs:
    steps:
```"""
    
    mock_response = Mock()
    mock_response.choices = [mock_choice]
    
    mock_client = Mock()
    mock_client.chat.completions.create.return_value = mock_response
    mock_openai.return_value = mock_client
    
    result = generate_ape_from_nl("Test prompt")
    
    assert result.startswith("task test:")
    assert "```" not in result


@patch('ape_openai.generator.OPENAI_AVAILABLE', True)
@patch('ape_openai.generator.OpenAI')
def test_generate_ape_from_nl_api_error(mock_openai):
    """Test handling of OpenAI API errors."""
    mock_client = Mock()
    mock_client.chat.completions.create.side_effect = Exception("API Error")
    mock_openai.return_value = mock_client
    
    with pytest.raises(Exception, match="Failed to generate Ape code"):
        generate_ape_from_nl("Test prompt")
