#!/usr/bin/env python3
"""
Tests for the Planning and Execution LLM Tool
"""

import os
import sys
import pytest
import asyncio
from pathlib import Path
from unittest.mock import Mock, patch, AsyncMock
from typing import Dict, Any

# Add the parent directory to the Python path
sys.path.append(str(Path(__file__).parent.parent))

from tools.plan_exec_llm import PlanningTool, PlanningResult

@pytest.fixture
def planning_tool():
    return PlanningTool()

@pytest.fixture
def mock_llm_response():
    return """
Technical feasibility: High
- Implementation is straightforward
- No major technical challenges

Implementation steps:
1. Step one
2. Step two
3. Step three

Success criteria:
- All tests pass
- Code coverage > 90%
"""

@pytest.mark.asyncio
async def test_plan_success(planning_tool, mock_llm_response):
    """Test successful plan generation."""
    with patch('tools.plan_exec_llm.query_llm', new_callable=AsyncMock) as mock_query:
        mock_query.return_value = mock_llm_response
        
        result = await planning_tool.plan("Test task")
        
        assert result.success is True
        assert result.plan == mock_llm_response
        assert result.error is None
        assert isinstance(result.context, dict)
        assert result.context["prompt"] == "Test task"
        assert result.context["file_content"] is None

@pytest.mark.asyncio
async def test_plan_with_file_content(planning_tool, mock_llm_response):
    """Test plan generation with file content."""
    file_content = "def test():\n    pass"
    
    with patch('tools.plan_exec_llm.query_llm', new_callable=AsyncMock) as mock_query:
        mock_query.return_value = mock_llm_response
        
        result = await planning_tool.plan("Test task", file_content)
        
        assert result.success is True
        assert result.plan == mock_llm_response
        assert result.error is None
        assert result.context["file_content"] == file_content

@pytest.mark.asyncio
async def test_plan_validation_failure(planning_tool):
    """Test plan validation failure."""
    invalid_response = "Invalid plan without required sections"
    
    with patch('tools.plan_exec_llm.query_llm', new_callable=AsyncMock) as mock_query:
        mock_query.return_value = invalid_response
        
        result = await planning_tool.plan("Test task")
        
        assert result.success is False
        assert result.plan is None
        assert "does not meet requirements" in result.error
        assert result.context["prompt"] == "Test task"

@pytest.mark.asyncio
async def test_plan_llm_error(planning_tool):
    """Test handling of LLM API error."""
    with patch('tools.plan_exec_llm.query_llm', new_callable=AsyncMock) as mock_query:
        mock_query.side_effect = Exception("API Error")
        
        result = await planning_tool.plan("Test task")
        
        assert result.success is False
        assert result.plan is None
        assert "API Error" in result.error
        assert result.context["prompt"] == "Test task"

@pytest.mark.asyncio
async def test_plan_caching(planning_tool, mock_llm_response):
    """Test plan caching functionality."""
    with patch('tools.plan_exec_llm.query_llm', new_callable=AsyncMock) as mock_query:
        mock_query.return_value = mock_llm_response
        
        # First call
        result1 = await planning_tool.plan("Test task")
        assert result1.success is True
        
        # Second call with same prompt should use cache
        result2 = await planning_tool.plan("Test task")
        assert result2.success is True
        assert result2.plan == result1.plan
        
        # LLM should only be called once
        assert mock_query.call_count == 1

def test_prepare_planning_prompt(planning_tool):
    """Test planning prompt preparation."""
    prompt = "Test prompt"
    result = planning_tool._prepare_planning_prompt(prompt)
    
    # Check for numbered items in the prompt
    assert "1. Technical feasibility" in result
    assert "4. Implementation steps" in result
    assert "5. Success criteria" in result
    assert prompt in result

def test_validate_plan(planning_tool):
    """Test plan validation."""
    valid_plan = """
    Technical feasibility: High
    Implementation steps: Step 1, Step 2
    Success criteria: All tests pass
    """
    assert planning_tool._validate_plan(valid_plan) is True
    
    invalid_plan = "Invalid plan"
    assert planning_tool._validate_plan(invalid_plan) is False

@pytest.mark.asyncio
async def test_plan_integration(planning_tool, mock_llm_response):
    """Test plan generation integration."""
    with patch('tools.plan_exec_llm.query_llm', new_callable=AsyncMock) as mock_query:
        mock_query.return_value = mock_llm_response
        
        result = await planning_tool.plan("Create a simple Python function to calculate factorial")
        assert result.success is True
        assert "Implementation steps" in result.plan

@pytest.mark.asyncio
async def test_error_handling(planning_tool):
    """Test error handling."""
    with patch('tools.plan_exec_llm.query_llm', new_callable=AsyncMock) as mock_query:
        mock_query.side_effect = Exception("Test error")
        
        result = await planning_tool.plan("")
        assert result.success is False
        assert "Test error" in result.error

@pytest.mark.asyncio
async def test_context_caching(planning_tool):
    """Test context caching."""
    prompt = "Test prompt"
    mock_response = "Technical feasibility: High\nImplementation steps: Step 1\nSuccess criteria: Pass"
    
    with patch('tools.plan_exec_llm.query_llm', new_callable=AsyncMock) as mock_query:
        mock_query.return_value = mock_response
        
        result1 = await planning_tool.plan(prompt)
        assert result1.success is True
        assert prompt in planning_tool.context_cache
        
        # Second call should use cache
        result2 = await planning_tool.plan(prompt)
        assert result2.plan == result1.plan
        assert mock_query.call_count == 1

def test_token_tracking(planning_tool):
    """Test token usage tracking."""
    assert hasattr(planning_tool, 'token_tracker')
    assert planning_tool.token_tracker is not None

if __name__ == "__main__":
    pytest.main([__file__]) 