#!/usr/bin/env python3
"""
Enhanced Perplexity API Module
Provides enhanced functionality for interacting with the Perplexity API.
"""

import os
import sys
import json
import logging
import asyncio
import aiohttp
from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from datetime import datetime
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class PerplexityResponse:
    """Response from Perplexity API."""
    text: str
    sources: List[Dict[str, str]]
    metadata: Dict[str, Any]
    timestamp: datetime

class EnhancedPerplexityAPI:
    """Enhanced interface for Perplexity API with additional features."""
    
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.getenv("PERPLEXITY_API_KEY")
        if not self.api_key:
            raise ValueError("Perplexity API key not found")
        
        self.base_url = "https://api.perplexity.ai/v1"
        self.session = None
        self.response_cache = {}
    
    async def __aenter__(self):
        """Set up async context."""
        self.session = aiohttp.ClientSession(
            headers={"Authorization": f"Bearer {self.api_key}"}
        )
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Clean up async context."""
        if self.session:
            await self.session.close()
    
    async def query(
        self,
        prompt: str,
        model: str = "pplx-7b-online",
        max_tokens: int = 1024,
        temperature: float = 0.7,
        top_p: float = 0.9,
        cache: bool = True
    ) -> PerplexityResponse:
        """
        Query the Perplexity API with enhanced functionality.
        
        Args:
            prompt: The query prompt
            model: Model to use (default: pplx-7b-online)
            max_tokens: Maximum tokens in response
            temperature: Sampling temperature
            top_p: Nucleus sampling parameter
            cache: Whether to use response caching
        
        Returns:
            PerplexityResponse object containing response and metadata
        """
        # Check cache
        cache_key = f"{prompt}:{model}:{max_tokens}:{temperature}:{top_p}"
        if cache and cache_key in self.response_cache:
            return self.response_cache[cache_key]
        
        if not self.session:
            self.session = aiohttp.ClientSession(
                headers={"Authorization": f"Bearer {self.api_key}"}
            )
        
        try:
            async with self.session.post(
                f"{self.base_url}/chat/completions",
                json={
                    "model": model,
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": max_tokens,
                    "temperature": temperature,
                    "top_p": top_p
                }
            ) as response:
                if response.status != 200:
                    error_text = await response.text()
                    raise Exception(f"API error: {error_text}")
                
                data = await response.json()
                
                result = PerplexityResponse(
                    text=data["choices"][0]["message"]["content"],
                    sources=self._extract_sources(data),
                    metadata={
                        "model": model,
                        "usage": data.get("usage", {}),
                        "finish_reason": data["choices"][0].get("finish_reason")
                    },
                    timestamp=datetime.now()
                )
                
                if cache:
                    self.response_cache[cache_key] = result
                
                return result
                
        except Exception as e:
            logger.error(f"Error querying Perplexity API: {str(e)}")
            raise
    
    async def query_with_context(
        self,
        prompt: str,
        context: str,
        **kwargs
    ) -> PerplexityResponse:
        """Query with additional context."""
        enhanced_prompt = f"Context:\n{context}\n\nQuery:\n{prompt}"
        return await self.query(enhanced_prompt, **kwargs)
    
    async def query_with_sources(
        self,
        prompt: str,
        sources: List[str],
        **kwargs
    ) -> PerplexityResponse:
        """Query with source documents."""
        sources_text = "\n".join(f"Source {i+1}:\n{source}" for i, source in enumerate(sources))
        enhanced_prompt = f"Sources:\n{sources_text}\n\nQuery:\n{prompt}"
        return await self.query(enhanced_prompt, **kwargs)
    
    def _extract_sources(self, response_data: Dict[str, Any]) -> List[Dict[str, str]]:
        """Extract source information from response."""
        sources = []
        try:
            message = response_data["choices"][0]["message"]
            if "context" in message:
                for source in message["context"].get("documents", []):
                    sources.append({
                        "title": source.get("title", ""),
                        "url": source.get("url", ""),
                        "snippet": source.get("snippet", "")
                    })
        except Exception as e:
            logger.warning(f"Error extracting sources: {str(e)}")
        return sources
    
    def clear_cache(self):
        """Clear the response cache."""
        self.response_cache.clear()
    
    async def batch_query(
        self,
        prompts: List[str],
        max_concurrent: int = 5,
        **kwargs
    ) -> List[PerplexityResponse]:
        """Execute multiple queries concurrently."""
        semaphore = asyncio.Semaphore(max_concurrent)
        
        async def _bounded_query(prompt: str) -> PerplexityResponse:
            async with semaphore:
                return await self.query(prompt, **kwargs)
        
        return await asyncio.gather(*[_bounded_query(prompt) for prompt in prompts])

async def main():
    """CLI interface for testing."""
    import argparse
    
    parser = argparse.ArgumentParser(description="Enhanced Perplexity API CLI")
    parser.add_argument("prompt", help="Query prompt")
    parser.add_argument("--model", default="pplx-7b-online", help="Model to use")
    parser.add_argument("--max-tokens", type=int, default=1024, help="Maximum tokens")
    parser.add_argument("--temperature", type=float, default=0.7, help="Temperature")
    parser.add_argument("--no-cache", action="store_true", help="Disable caching")
    args = parser.parse_args()
    
    try:
        async with EnhancedPerplexityAPI() as api:
            response = await api.query(
                args.prompt,
                model=args.model,
                max_tokens=args.max_tokens,
                temperature=args.temperature,
                cache=not args.no_cache
            )
            
            print("\nResponse:")
            print(response.text)
            
            if response.sources:
                print("\nSources:")
                for source in response.sources:
                    print(f"- {source['title']} ({source['url']})")
            
            print("\nMetadata:")
            print(json.dumps(response.metadata, indent=2))
            
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(main()) 