"""OpenRouter research service implementation."""

import asyncio
import logging
import time
from typing import List, Optional
from datetime import datetime, timezone

# Import config - handle both direct and package imports
try:
    from ...config import config
except ImportError:
    # Handle when run directly
    import sys
    from pathlib import Path
    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
    from config import config

# Import protocols with fallback for dataclasses
try:
    from ..protocols import ResearchServiceProtocol, ResearchQuery, ResearchResult
except ImportError:
    # Create simple dataclasses if protocols import fails
    from dataclasses import dataclass
    from typing import List, Dict, Optional, Protocol, runtime_checkable
    
    @dataclass
    class ResearchQuery:
        query: str
        context: Optional[str] = None
        max_tokens: int = 4000
        temperature: float = 0.1
        model: str = "perplexity/sonar-reasoning-pro"
    
    @dataclass  
    class ResearchResult:
        query: str
        answer: str
        sources: List[str]
        confidence_score: float
        processing_time: float
        model_used: str
        token_usage: Dict[str, int]
    
    @runtime_checkable
    class ResearchServiceProtocol(Protocol):
        async def research(self, query: ResearchQuery) -> ResearchResult: ...
        async def batch_research(self, queries: List[ResearchQuery]) -> List[ResearchResult]: ...
        def get_available_models(self) -> List[str]: ...

# Import dependencies with fallback for standalone execution
try:
    from .api import OpenRouterAPIClient
    from .config import ResearchConfig
    from .errors import ResearchError, ResearchConfigurationError
except ImportError:
    # Handle when run directly - import from same directory
    import sys
    from pathlib import Path
    current_dir = Path(__file__).parent
    sys.path.insert(0, str(current_dir))
    
    from api import OpenRouterAPIClient
    from config import ResearchConfig
    from errors import ResearchError, ResearchConfigurationError

logger = logging.getLogger(__name__)


class OpenRouterResearchService:
    """OpenRouter AI research service for knowledge queries."""

    def __init__(self, research_config: Optional[ResearchConfig] = None):
        """Initialize research service.
        
        Args:
            research_config: Research service configuration
            
        """
        self.name = "OpenRouterResearchService"
        
        # Use provided config or create from environment
        if research_config:
            self.config = research_config
        else:
            if not config.OPENROUTER_API_KEY:
                raise ResearchConfigurationError("OPENROUTER_API_KEY not found in environment")
            self.config = ResearchConfig.from_env(config.OPENROUTER_API_KEY)
        
        # Validate configuration
        self.config.validate()
        
        # Initialize API client
        self.api_client = OpenRouterAPIClient(self.config)
        
        logger.info("OpenRouter research service initialized")

    async def research(self, query: ResearchQuery) -> ResearchResult:
        """Perform a research query.
        
        Args:
            query: Research query to execute
            
        Returns:
            Research result with answer and sources
            
        Raises:
            ResearchError: If research fails
            
        """
        start_time = time.time()
        
        try:
            logger.info(f"Starting research query: {query.query[:100]}...")
            
            # Perform the API call
            response_data = await self.api_client.research_query(
                query=query.query,
                context=query.context,
                model=query.model,
                max_tokens=query.max_tokens,
                temperature=query.temperature
            )
            
            # Extract response content
            answer = self.api_client.extract_response_content(response_data)
            token_usage = self.api_client.extract_token_usage(response_data)
            
            # Extract sources from the response (Perplexity models often include citations)
            sources = self._extract_sources(answer)
            
            # Calculate confidence score (simple heuristic based on response length and sources)
            confidence_score = self._calculate_confidence(answer, sources, token_usage)
            
            processing_time = time.time() - start_time
            
            result = ResearchResult(
                query=query.query,
                answer=answer,
                sources=sources,
                confidence_score=confidence_score,
                processing_time=processing_time,
                model_used=response_data["_metadata"]["model_used"],
                token_usage=token_usage
            )
            
            logger.info(f"Research completed in {processing_time:.2f}s with {len(sources)} sources")
            return result
            
        except Exception as e:
            processing_time = time.time() - start_time
            logger.error(f"Research failed after {processing_time:.2f}s: {e}")
            raise ResearchError(f"Research query failed: {e}")

    async def batch_research(self, queries: List[ResearchQuery]) -> List[ResearchResult]:
        """Perform multiple research queries.
        
        Args:
            queries: List of research queries to execute
            
        Returns:
            List of research results
            
        """
        logger.info(f"Starting batch research with {len(queries)} queries")
        
        # Execute queries concurrently with reasonable concurrency limit
        semaphore = asyncio.Semaphore(3)  # Limit to 3 concurrent requests
        
        async def research_with_semaphore(query: ResearchQuery) -> ResearchResult:
            async with semaphore:
                return await self.research(query)
        
        results = await asyncio.gather(
            *[research_with_semaphore(query) for query in queries],
            return_exceptions=True
        )
        
        # Convert exceptions to error results
        processed_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                error_result = ResearchResult(
                    query=queries[i].query,
                    answer=f"Error: {str(result)}",
                    sources=[],
                    confidence_score=0.0,
                    processing_time=0.0,
                    model_used=queries[i].model,
                    token_usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
                )
                processed_results.append(error_result)
            else:
                processed_results.append(result)
        
        logger.info(f"Batch research completed with {len(processed_results)} results")
        return processed_results

    def get_available_models(self) -> List[str]:
        """Get list of available research models.
        
        Returns:
            List of available model names
            
        """
        # Return commonly available Perplexity models through OpenRouter
        return [
            "perplexity/sonar-reasoning-pro",
            "perplexity/sonar-reasoning", 
            "perplexity/sonar-pro",
            "perplexity/sonar",
            "openai/gpt-4o",
            "openai/gpt-4o-mini",
            "anthropic/claude-3.5-sonnet",
            "anthropic/claude-3.5-haiku"
        ]

    def _extract_sources(self, answer: str) -> List[str]:
        """Extract source URLs from research answer.
        
        Args:
            answer: Research answer text
            
        Returns:
            List of source URLs
            
        """
        import re
        
        # Look for common URL patterns in the answer
        url_pattern = r'https?://[^\s\])\}]+'
        sources = re.findall(url_pattern, answer)
        
        # Clean up and deduplicate sources
        cleaned_sources = []
        seen = set()
        
        for source in sources:
            # Remove trailing punctuation
            source = source.rstrip('.,;:!?')
            if source not in seen and len(source) > 10:
                cleaned_sources.append(source)
                seen.add(source)
        
        return cleaned_sources[:self.config.max_sources]

    def _calculate_confidence(self, answer: str, sources: List[str], token_usage: dict) -> float:
        """Calculate confidence score for research result.
        
        Args:
            answer: Research answer text
            sources: List of source URLs
            token_usage: Token usage information
            
        Returns:
            Confidence score between 0.0 and 1.0
            
        """
        # Base score from answer length and structure
        answer_length = len(answer.strip())
        if answer_length < 50:
            length_score = 0.3
        elif answer_length < 200:
            length_score = 0.6
        elif answer_length < 500:
            length_score = 0.8
        else:
            length_score = 0.9
        
        # Bonus for having sources
        source_score = min(len(sources) * 0.1, 0.3)
        
        # Bonus for specific indicators
        indicators = ['according to', 'research shows', 'studies indicate', 'data suggests']
        indicator_score = sum(0.05 for indicator in indicators if indicator.lower() in answer.lower())
        indicator_score = min(indicator_score, 0.2)
        
        # Combine scores
        confidence = length_score + source_score + indicator_score
        return min(confidence, 1.0)
