import asyncio
from typing import Dict, Any

import httpx
from pydantic import BaseModel

from model_registry import model_card

perpelxity_card = model_card["perplexity"]


class QueryRequest(BaseModel):
    query: str


class PerplexityClient:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = perpelxity_card["url"]
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

    async def generate_response(self, query: str) -> Dict[str, Any]:
        payload = {
            "model": perpelxity_card["model"],
            "messages": [
                {
                    "role": "system",
                    "content": perpelxity_card["prompt"]
                },
                {
                    "role": "user",
                    "content": query
                }
            ],
            "max_tokens": perpelxity_card["inference_config"]["max_tokens"],
            "temperature": perpelxity_card["inference_config"]["temperature"],
            "stream": False,  # We'll handle streaming separately
        }
        
        async with httpx.AsyncClient() as client:
            try:
                async with httpx.AsyncClient(
                    timeout=httpx.Timeout(
                        connect=10.0,    # Connection timeout
                        read=45.0,       # Read timeout
                        write=10.0,      # Write timeout
                        pool=10.0        # Connection pool timeout
                    )
                ) as client:
                    response = await client.post(
                        self.base_url,
                        headers=self.headers,
                        json=payload
                    )
                    response.raise_for_status()
                    return response.json()
            except httpx.HTTPStatusError as e:
                print(f"HTTP error occurred: {e}")
                print(f"Response text: {e.response.text}")
                raise
            except httpx.RequestError as e:
                print(f"Request error occurred: {e}")
                raise
            except Exception as e:
                print(f"An unexpected error occurred: {e}")
                raise


def parse_perplexity_response(response: Dict[str, Any]) -> Dict[str, Any]:
    """
    Parse the Perplexity API response and extract key information.
    
    Args:
        response (Dict[str, Any]): Raw response from Perplexity API
    
    Returns:
        Dict[str, Any]: Parsed response with content and citations
    """
    print("parse_perplexity_response called...")
    # Basic citation extraction (this is a simple implementation)
    # In a real-world scenario, you might want a more sophisticated citation extraction
    citations = []
    
    # default content to stream if no response from the Perplexity AI
    default_content = (
        "I'm sorry, I couldn't find any relevant information for your query from the available sources. "
        "If you'd like, you can try rephrasing your question or provide more context to help refine the search. "
        "Alternatively, let me know if you'd like assistance in a different area."
        )
    
    citations = response.get("citations", [])
    content = response.get("choices", default_content)[0]["message"]["content"]
    return {
        "content": content,
        "citations": citations
    }