from dataclasses import dataclass, field
from os import getenv
from typing import Any, Dict, Optional, Type, Union

from pydantic import BaseModel

from agno.exceptions import ModelProviderError
from agno.models.message import Citations, UrlCitation
from agno.models.metrics import Metrics
from agno.models.response import ModelResponse
from agno.utils.log import log_debug, log_warning

try:
    from openai.types.chat.chat_completion import ChatCompletion
    from openai.types.chat.chat_completion_chunk import (
        ChatCompletionChunk,
        ChoiceDelta,
    )
    from openai.types.chat.parsed_chat_completion import ParsedChatCompletion
    from openai.types.completion_usage import CompletionUsage
except ModuleNotFoundError:
    raise ImportError("`openai` not installed. Please install using `pip install openai`")

from agno.models.openai.like import OpenAILike


@dataclass
class Perplexity(OpenAILike):
    """
    A class for using models hosted on Perplexity.

    Attributes:
        id (str): The model id. Defaults to "sonar".
        name (str): The model name. Defaults to "Perplexity".
        provider (str): The provider name. Defaults to "Perplexity".
        api_key (Optional[str]): The API key.
        base_url (str): The base URL. Defaults to "https://api.perplexity.ai/chat/completions".
        max_tokens (int): The maximum number of tokens. Defaults to 1024.
    """

    id: str = "sonar"
    name: str = "Perplexity"
    provider: str = "Perplexity"

    api_key: Optional[str] = field(default_factory=lambda: getenv("PERPLEXITY_API_KEY"))
    base_url: str = "https://api.perplexity.ai/"
    max_tokens: int = 1024
    top_k: Optional[float] = None

    supports_native_structured_outputs: bool = False
    supports_json_schema_outputs: bool = True

    def get_request_params(
        self,
        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """
        Returns keyword arguments for API requests.
        """
        # Define base request parameters
        base_params: Dict[str, Any] = {
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "presence_penalty": self.presence_penalty,
            "frequency_penalty": self.frequency_penalty,
        }

        if response_format is not None:
            base_params["response_format"] = response_format

        # Filter out None values
        request_params = {k: v for k, v in base_params.items() if v is not None}
        # Add additional request params if provided
        if self.request_params:
            request_params.update(self.request_params)

        if request_params:
            log_debug(f"Calling {self.provider} with request parameters: {request_params}", log_level=2)
        return request_params

    def parse_provider_response(self, response: Union[ChatCompletion, ParsedChatCompletion], **kwargs) -> ModelResponse:
        """
        Parse the Perplexity response into a ModelResponse.

        Args:
            response: Response from invoke() method

        Returns:
            ModelResponse: Parsed response data
        """
        model_response = ModelResponse()

        if hasattr(response, "error") and response.error:
            raise ModelProviderError(
                message=response.error.get("message", "Unknown model error"),
                model_name=self.name,
                model_id=self.id,
            )

        # Get response message
        response_message = response.choices[0].message

        # Add role
        if response_message.role is not None:
            model_response.role = response_message.role

        # Add content
        if response_message.content is not None:
            model_response.content = response_message.content

        # Add tool calls
        if response_message.tool_calls is not None and len(response_message.tool_calls) > 0:
            try:
                model_response.tool_calls = [t.model_dump() for t in response_message.tool_calls]
            except Exception as e:
                log_warning(f"Error processing tool calls: {e}")

        # Add citations if present
        if hasattr(response, "citations") and response.citations is not None:
            model_response.citations = Citations(
                urls=[UrlCitation(url=c) for c in response.citations],
            )

        if response.usage is not None:
            model_response.response_usage = self._get_metrics(response.usage)

        return model_response

    def parse_provider_response_delta(self, response_delta: ChatCompletionChunk) -> ModelResponse:
        """
        Parse the Perplexity streaming response into a ModelResponse.

        Args:
            response_delta: Raw response chunk from Perplexity

        Returns:
            ProviderResponse: Iterator of parsed response data
        """
        model_response = ModelResponse()
        if response_delta.choices and len(response_delta.choices) > 0:
            choice_delta: ChoiceDelta = response_delta.choices[0].delta

            if choice_delta:
                # Add content
                if choice_delta.content is not None:
                    model_response.content = choice_delta.content

                # Add tool calls
                if choice_delta.tool_calls is not None:
                    model_response.tool_calls = choice_delta.tool_calls  # type: ignore

        # Add citations if present
        if hasattr(response_delta, "citations") and response_delta.citations is not None:
            model_response.citations = Citations(
                urls=[UrlCitation(url=c) for c in response_delta.citations],
            )

        # Add usage metrics if present
        if response_delta.usage is not None:
            model_response.response_usage = self._get_metrics(response_delta.usage)

        return model_response

    def _get_metrics(self, response_usage: CompletionUsage) -> Metrics:
        """
        Parse the given Perplexity usage into an Agno Metrics object.
        """
        metrics = Metrics()

        metrics.input_tokens = response_usage.prompt_tokens or 0
        metrics.output_tokens = response_usage.completion_tokens or 0
        metrics.total_tokens = response_usage.total_tokens or 0

        # Add the prompt_tokens_details field
        if prompt_token_details := response_usage.prompt_tokens_details:
            metrics.audio_input_tokens = prompt_token_details.audio_tokens or 0
            metrics.cache_read_tokens = prompt_token_details.cached_tokens or 0

        # Add the completion_tokens_details field
        if completion_tokens_details := response_usage.completion_tokens_details:
            metrics.audio_output_tokens = completion_tokens_details.audio_tokens or 0
            metrics.reasoning_tokens = completion_tokens_details.reasoning_tokens or 0

        return metrics
