from typing import Callable, Any, List, Dict, Union, Optional

import requests
from pydantic import BaseModel
from dotenv import load_dotenv

import os

load_dotenv()


class TrimRouterResponse(BaseModel):
    chosen_model_id: str
    request_id: str
    prompt: Optional[str] = None
    router_reasoning: Optional[str] = None


class Trim:
    def __init__(self, api_key: Optional[str] = None):
        """
        Initialize Trim Client.

        Args:
            api_key (str, optional): API key for authentication
        """
        self.base_url = "https://llm-gateway.net/api"
        self.api_key = api_key or os.getenv("TRIM_API_KEY")
        self.session = requests.Session()
        if not self.api_key:
            raise Exception(
                "The api_key client option must be set by setting the TRIM_API_KEY environment variable")

        self.session.headers.update({"Authorization": f"Bearer {self.api_key}"})

    def route_llm_request(
            self,
            prompt: str,
            intended_model_id: Optional[str] = None
    ) -> TrimRouterResponse:
        """
        Route an LLM request to determine the appropriate model.

        Args:
            prompt (str): The prompt to route
            intended_model_id (str, optional): The intended model ID to use

        Returns:
            TrimRouterResponse: The routing response containing chosen model and request ID
        """
        url = f"{self.base_url}/llm/route"
        headers = {"Content-Type": "application/json"}
        payload = {
            "prompt": prompt
        }
        if intended_model_id:
            payload["intended_model_id"] = intended_model_id

        response = self.session.post(url, headers=headers, json=payload)
        response.raise_for_status()

        return TrimRouterResponse(**response.json())

    def cache_llm_response(
            self,
            request_id: str,
            response: str,
            number_of_input_tokens: Optional[int] = None,
            number_of_output_tokens: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        Cache an LLM response.

        Args:
            request_id (str): The request ID to cache the response for
            response (str): The response to cache
            number_of_input_tokens (int, optional): Number of tokens in the input
            number_of_output_tokens (int, optional): Number of tokens in the output

        Returns:
            Dict[str, Any]: The server response
        """
        url = f"{self.base_url}/llm/cache"

        payload = {
            "request_id": request_id,
            "response": response
        }

        if number_of_input_tokens is not None:
            payload["number_of_input_tokens"] = number_of_input_tokens
        if number_of_output_tokens is not None:
            payload["number_of_output_tokens"] = number_of_output_tokens

        response = self.session.post(url, json=payload)
        response.raise_for_status()

        return response.json()


class TrimOpenAICompletions:
    """
    A wrapper class that routes OpenAI-style chat completion requests
    through a custom LLM routing client to select an optimal model.

    Attributes:
        original_create_function (Callable): The original function for creating completions.
        trim_client (Trim): Client responsible for routing logic.
    """

    def __init__(self, original_create_function: Callable):
        """
        Initializes with the original OpenAI completion function.

        Args:
            original_create_function (Callable): The function used to generate completions.
        """
        self.original_create_function: Callable = original_create_function
        self.trim_client: Trim = Trim()

    def __call__(self, model: str, messages: List[Dict[str, Union[str, Any]]], **kwargs) -> Any:
        """
        Routes the chat completion request to an appropriate model.

        Args:
            model (str): Intended model name.
            messages (List[Dict[str, Union[str, Any]]]): List of chat message dicts.
            **kwargs: Additional parameters for the completion function.

        Returns:
            Any: Result from the chosen model's completion function.
        """
        routing_response = self.trim_client.route_llm_request(
            prompt=str(messages),
            intended_model_id=model
        )

        results = self.original_create_function(
            model=routing_response.chosen_model_id,
            messages=messages,
            **kwargs
        )

        response = results.choices[-1].message.content
        n_input_tokens = results.usage.prompt_tokens
        n_output_tokens = results.usage.completion_tokens

        cache_response = self.trim_client.cache_llm_response(
            request_id=routing_response.request_id,
            response=response,
            number_of_input_tokens=n_input_tokens,
            number_of_output_tokens=n_output_tokens
        )

        return results


class TrimOpenAIResponses:
    """
    A wrapper class that routes OpenAI-style response generation requests
    through a custom LLM routing client to select an optimal model.

    Attributes:
        original_create_function (Callable): The original function for generating responses.
        trim_client (Trim): Client responsible for routing logic.
    """

    def __init__(self, original_create_function: Callable):
        """
        Initializes the original OpenAI response function.

        Args:
            original_create_function (Callable): The function used to generate responses.
        """
        self.original_create_function: Callable = original_create_function
        self.trim_client: Trim = Trim()

    def __call__(self, model: str, input: Union[str, Dict[str, Any]], **kwargs) -> Any:
        """
        Routes the response generation request to an appropriate model.

        Args:
            model (str): Intended model name.
            input (Union[str, Dict[str, Any]]): Prompt or structured input.
            **kwargs: Additional parameters for the response function.

        Returns:
            Any: Result from the chosen model's response function.
        """
        routing_response = self.trim_client.route_llm_request(
            prompt=str(input),
            intended_model_id=model
        )

        results = self.original_create_function(
            model=routing_response.chosen_model_id,
            input=input,
            **kwargs
        )

        response = results.output[-1].content[-1].text
        n_input_tokens = results.usage.input_tokens
        n_output_tokens = results.usage.output_tokens

        cache_response = self.trim_client.cache_llm_response(
            request_id=routing_response.request_id,
            response=response,
            number_of_input_tokens=n_input_tokens,
            number_of_output_tokens=n_output_tokens
        )

        return results
