"""OpenRouter API client for research service."""

import asyncio
import json
import logging
import time
from typing import Any, Dict, List
from datetime import datetime, timezone

import httpx
from tenacity import AsyncRetrying, stop_after_attempt, wait_exponential, retry_if_exception_type

# Import dependencies with fallback for standalone execution
try:
    from .config import ResearchConfig
    from .errors import ResearchAPIError, ResearchTimeoutError, ResearchRateLimitError
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 config import ResearchConfig
    from errors import ResearchAPIError, ResearchTimeoutError, ResearchRateLimitError

logger = logging.getLogger(__name__)


class OpenRouterAPIClient:
    """OpenRouter API client for research queries."""

    def __init__(self, config: ResearchConfig):
        """Initialize OpenRouter API client.

        Args:
            config: Research service configuration

        """
        self.config = config
        self.base_url = config.base_url
        self.api_key = config.api_key
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://trax.local",  # Optional: your site URL
            "X-Title": "Trax Research Agent",  # Optional: your app name
        }

    async def research_query(
        self, 
        query: str, 
        context: str = None,
        model: str = None,
        max_tokens: int = None,
        temperature: float = None,
        **kwargs
    ) -> Dict[str, Any]:
        """Perform a research query using OpenRouter.

        Args:
            query: Research question/prompt
            context: Additional context for the query
            model: Model to use (defaults to config default)
            max_tokens: Maximum tokens in response
            temperature: Sampling temperature
            **kwargs: Additional API parameters

        Returns:
            Response data from OpenRouter API

        Raises:
            ResearchAPIError: If API call fails
            ResearchTimeoutError: If request times out
            ResearchRateLimitError: If rate limit exceeded

        """
        # Build the full prompt
        if context:
            full_prompt = f"Context: {context}\n\nQuery: {query}"
        else:
            full_prompt = query

        # Prepare request data
        request_data = {
            "model": model or self.config.default_model,
            "messages": [
                {
                    "role": "user",
                    "content": full_prompt
                }
            ],
            "max_tokens": max_tokens or self.config.max_tokens,
            "temperature": temperature or self.config.temperature,
            "stream": False
        }

        # Add any additional parameters
        request_data.update(kwargs)

        # Perform request with retries
        async for attempt in AsyncRetrying(
            stop=stop_after_attempt(self.config.max_retries),
            wait=wait_exponential(multiplier=1, min=4, max=10),
            retry=retry_if_exception_type((ResearchAPIError, httpx.RequestError))
        ):
            with attempt:
                try:
                    async with httpx.AsyncClient(
                        timeout=httpx.Timeout(self.config.timeout)
                    ) as client:
                        start_time = time.time()
                        
                        response = await client.post(
                            f"{self.base_url}/chat/completions",
                            headers=self.headers,
                            json=request_data
                        )
                        
                        processing_time = time.time() - start_time
                        
                        # Handle rate limiting
                        if response.status_code == 429:
                            retry_after = int(response.headers.get("Retry-After", 60))
                            raise ResearchRateLimitError(
                                f"Rate limit exceeded. Retry after {retry_after} seconds.",
                                retry_after=retry_after
                            )
                        
                        # Handle other HTTP errors
                        if response.status_code != 200:
                            error_data = {}
                            try:
                                error_data = response.json()
                            except:
                                pass
                            
                            raise ResearchAPIError(
                                f"API request failed: {response.status_code} - {response.text}",
                                status_code=response.status_code,
                                response_data=error_data
                            )
                        
                        response_data = response.json()
                        
                        # Add metadata
                        response_data["_metadata"] = {
                            "processing_time": processing_time,
                            "timestamp": datetime.now(timezone.utc).isoformat(),
                            "model_used": request_data["model"],
                            "query": query,
                            "context": context
                        }
                        
                        logger.info(f"Research query completed in {processing_time:.2f}s")
                        return response_data

                except httpx.TimeoutException as e:
                    raise ResearchTimeoutError(f"Request timed out: {e}")
                except httpx.RequestError as e:
                    logger.warning(f"Request error (attempt {attempt.retry_state.attempt_number}): {e}")
                    raise ResearchAPIError(f"Request error: {e}")

    async def get_available_models(self) -> List[str]:
        """Get list of available models from OpenRouter.

        Returns:
            List of available model names

        """
        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(30)) as client:
                response = await client.get(
                    f"{self.base_url}/models",
                    headers=self.headers
                )
                
                if response.status_code == 200:
                    models_data = response.json()
                    return [model["id"] for model in models_data.get("data", [])]
                else:
                    logger.warning(f"Failed to fetch models: {response.status_code}")
                    return [self.config.default_model]
                    
        except Exception as e:
            logger.warning(f"Error fetching available models: {e}")
            return [self.config.default_model]

    def extract_response_content(self, response_data: Dict[str, Any]) -> str:
        """Extract the main content from OpenRouter response.

        Args:
            response_data: Response from OpenRouter API

        Returns:
            Extracted text content

        """
        try:
            return response_data["choices"][0]["message"]["content"]
        except (KeyError, IndexError) as e:
            raise ResearchAPIError(f"Invalid response format: {e}")

    def extract_token_usage(self, response_data: Dict[str, Any]) -> Dict[str, int]:
        """Extract token usage from OpenRouter response.

        Args:
            response_data: Response from OpenRouter API

        Returns:
            Token usage information

        """
        try:
            usage = response_data.get("usage", {})
            return {
                "prompt_tokens": usage.get("prompt_tokens", 0),
                "completion_tokens": usage.get("completion_tokens", 0),
                "total_tokens": usage.get("total_tokens", 0)
            }
        except Exception:
            return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
