#!/usr/bin/env python3
"""
Tests for the Enhanced Perplexity API module
"""

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

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

from bat_genie.perplexity_enhanced import EnhancedPerplexityAPI, PerplexityResponse

@pytest.fixture
async def mock_response_data():
    return {
        "choices": [{
            "message": {
                "content": "Test response",
                "context": {
                    "documents": [
                        {
                            "title": "Test Document",
                            "url": "https://example.com",
                            "snippet": "Test snippet"
                        }
                    ]
                }
            },
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": 10,
            "completion_tokens": 20,
            "total_tokens": 30
        }
    }

@pytest.fixture
async def mock_api_response(mock_response_data):
    mock = AsyncMock()
    mock.__aenter__.return_value.status = 200
    mock.__aenter__.return_value.json.return_value = mock_response_data
    return mock

@pytest.fixture
async def api(mock_api_response):
    with patch('aiohttp.ClientSession.post', return_value=mock_api_response):
        api = EnhancedPerplexityAPI("test_api_key")
        return api

@pytest.mark.asyncio
async def test_initialization():
    """Test API initialization."""
    api = EnhancedPerplexityAPI("test_api_key")
    assert api.api_key == "test_api_key"
    assert api.base_url == "https://api.perplexity.ai"

@pytest.mark.asyncio
async def test_query(api, mock_api_response):
    """Test basic query functionality."""
    response = await api.query("Test prompt")
    assert isinstance(response, PerplexityResponse)
    assert response.content == "Test response"

@pytest.mark.asyncio
async def test_query_with_context(api, mock_api_response):
    """Test querying with additional context."""
    response = await api.query_with_context(
        prompt="Test prompt",
        context="Additional context"
    )
    assert isinstance(response, PerplexityResponse)
    assert response.content == "Test response"

@pytest.mark.asyncio
async def test_query_with_sources(api, mock_api_response):
    """Test querying with source documents."""
    response = await api.query_with_sources(
        prompt="Test prompt",
        sources=["doc1.txt", "doc2.txt"]
    )
    assert isinstance(response, PerplexityResponse)
    assert response.content == "Test response"

@pytest.mark.asyncio
async def test_caching(api, mock_api_response):
    """Test response caching functionality."""
    response1 = await api.query("Test prompt")
    response2 = await api.query("Test prompt")
    assert response1 == response2

@pytest.mark.asyncio
async def test_batch_query(api, mock_api_response):
    """Test batch query functionality."""
    prompts = ["Prompt 1", "Prompt 2", "Prompt 3"]
    responses = await api.batch_query(prompts, max_concurrent=2)
    assert len(responses) == len(prompts)
    for response in responses:
        assert isinstance(response, PerplexityResponse)

@pytest.mark.asyncio
async def test_error_handling(api):
    """Test error handling."""
    with patch('aiohttp.ClientSession.post', side_effect=Exception("API error")):
        with pytest.raises(Exception) as exc_info:
            await api.query("Test prompt")
        assert "API error" in str(exc_info.value)

@pytest.mark.asyncio
async def test_source_extraction(api, mock_response_data):
    """Test source extraction from response."""
    sources = api._extract_sources(mock_response_data)
    assert len(sources) == 1
    assert sources[0]["title"] == "Test Document"
    assert sources[0]["url"] == "https://example.com"

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