"""Unit tests for the Perplexity research agent using OpenRouter."""

import pytest
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
from datetime import datetime, timezone
from typing import Dict, Any

from src.services.protocols import ResearchQuery, ResearchResult
from src.services.research.service import OpenRouterResearchService
from src.services.research.config import ResearchConfig


class TestPerplexityResearchAgent:
    """Test suite for Perplexity sonar-reasoning-pro research agent."""

    @pytest.fixture
    def mock_openrouter_service(self):
        """Create a mock OpenRouter service."""
        service = AsyncMock(spec=OpenRouterResearchService)
        service.research = AsyncMock()
        service.get_available_models = MagicMock(return_value=[
            "perplexity/sonar-reasoning-pro",
            "perplexity/sonar-medium-online",
            "openai/gpt-4"
        ])
        return service

    @pytest.fixture
    def sample_sonar_query(self):
        """Create a sample query for sonar-reasoning-pro."""
        return ResearchQuery(
            query="What are the latest developments in AI reasoning models?",
            context="Focus on models like o1, o3, and reasoning capabilities",
            max_tokens=4000,
            temperature=0.1,
            model="perplexity/sonar-reasoning-pro"
        )

    @pytest.fixture
    def sample_sonar_result(self):
        """Create a sample result from sonar-reasoning-pro."""
        return ResearchResult(
            query="What are the latest developments in AI reasoning models?",
            answer="Recent developments include OpenAI's o1 model which demonstrates...",
            sources=[
                "https://arxiv.org/abs/2024.12345",
                "https://openai.com/blog/o1-reasoning"
            ],
            confidence_score=0.92,
            processing_time=3.2,
            model_used="perplexity/sonar-reasoning-pro",
            token_usage={
                "prompt_tokens": 200,
                "completion_tokens": 1200,
                "total_tokens": 1400
            }
        )

    @pytest.mark.asyncio
    async def test_sonar_research_execution(self, mock_openrouter_service, sample_sonar_query, sample_sonar_result):
        """Test sonar-reasoning-pro research execution."""
        mock_openrouter_service.research.return_value = sample_sonar_result
        
        result = await mock_openrouter_service.research(sample_sonar_query)
        
        mock_openrouter_service.research.assert_called_once_with(sample_sonar_query)
        assert result.model_used == "perplexity/sonar-reasoning-pro"
        assert result.confidence_score > 0.9  # Sonar should have high confidence
        assert len(result.sources) > 0

    def test_sonar_query_validation(self):
        """Test sonar-specific query validation."""
        query = ResearchQuery(
            query="Test reasoning query",
            model="perplexity/sonar-reasoning-pro",
            max_tokens=4000,
            temperature=0.1
        )
        
        assert query.model == "perplexity/sonar-reasoning-pro"
        assert query.max_tokens <= 4000  # Sonar limit
        assert 0.0 <= query.temperature <= 1.0

    def test_sonar_result_quality(self, sample_sonar_result):
        """Test sonar result quality metrics."""
        assert sample_sonar_result.confidence_score >= 0.8  # High confidence expected
        assert sample_sonar_result.processing_time < 10.0  # Reasonable speed
        assert len(sample_sonar_result.sources) >= 1  # Should have sources
        assert sample_sonar_result.token_usage["total_tokens"] > 0

    @pytest.mark.asyncio
    async def test_sonar_error_handling(self, mock_openrouter_service, sample_sonar_query):
        """Test sonar error handling."""
        mock_openrouter_service.research.side_effect = Exception("OpenRouter API Error")
        
        with pytest.raises(Exception, match="OpenRouter API Error"):
            await mock_openrouter_service.research(sample_sonar_query)

    def test_sonar_model_availability(self, mock_openrouter_service):
        """Test sonar model availability."""
        models = mock_openrouter_service.get_available_models()
        
        assert "perplexity/sonar-reasoning-pro" in models
        assert "perplexity/sonar-medium-online" in models

    @pytest.mark.asyncio
    async def test_sonar_performance_benchmark(self, mock_openrouter_service, sample_sonar_query):
        """Test sonar performance benchmarks."""
        result = ResearchResult(
            query=sample_sonar_query.query,
            answer="Comprehensive reasoning analysis...",
            sources=["https://example.com"],
            confidence_score=0.95,
            processing_time=2.8,
            model_used="perplexity/sonar-reasoning-pro",
            token_usage={"total_tokens": 1500}
        )
        
        mock_openrouter_service.research.return_value = result
        
        research_result = await mock_openrouter_service.research(sample_sonar_query)
        
        # Performance expectations for sonar-reasoning-pro
        assert research_result.processing_time < 5.0  # Should be fast
        assert research_result.confidence_score >= 0.8  # High confidence
        assert research_result.token_usage["total_tokens"] < 2000  # Reasonable token usage


class TestOpenRouterIntegration:
    """Integration tests for OpenRouter with Perplexity."""

    @pytest.fixture
    def openrouter_config(self):
        """Create OpenRouter configuration."""
        return ResearchConfig(
            api_key="test-openrouter-key",
            base_url="https://openrouter.ai/api/v1",
            default_model="perplexity/sonar-reasoning-pro",
            max_tokens=4000,
            temperature=0.1
        )

    def test_openrouter_config_validation(self, openrouter_config):
        """Test OpenRouter configuration."""
        assert openrouter_config.api_key == "test-openrouter-key"
        assert openrouter_config.default_model == "perplexity/sonar-reasoning-pro"
        openrouter_config.validate()

    @pytest.mark.asyncio
    async def test_openrouter_service_initialization(self, openrouter_config):
        """Test OpenRouter service initialization."""
        service = OpenRouterResearchService(openrouter_config)
        
        assert hasattr(service, 'research')
        assert hasattr(service, 'get_available_models')
        assert callable(service.research)
        assert callable(service.get_available_models)


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