import requests
from typing import List, Dict, Any, Optional, Union
from .exceptions import UnauthorizedError, NotFoundError


class ModelsManagementClient:
    def __init__(self, base_url: str, api_key: Optional[str] = None):
        """
        Initialize the ModelsManagementClient.

        Args:
            base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
            api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
        """
        self._base_url = base_url.rstrip("/")  # Remove trailing slash if present
        self._api_key = api_key

    def _get_headers(self) -> Dict[str, str]:
        """
        Get the headers for API requests, including authorization if api_key is set.

        Returns:
            Dict[str, str]: Headers to use for API requests
        """
        headers = {}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        return headers

    def list(self, return_request: bool = False) -> Union[List[Dict[str, Any]], requests.Request]:
        """
        Get the list of models supported by the server.

        Args:
            return_request (bool): If True, returns the prepared request object instead of executing it.
                                 Useful for inspection or modification before sending.

        Returns:
            Union[List[Dict[str, Any]], requests.Request]: Either a list of model information dictionaries
            or a prepared request object if return_request is True.

        Raises:
            UnauthorizedError: If the request fails with a 401 status code
            requests.exceptions.RequestException: If the request fails with any other error
        """
        url = f"{self._base_url}/models"
        request = requests.Request("GET", url, headers=self._get_headers())

        if return_request:
            return request

        # Prepare and send the request
        session = requests.Session()
        try:
            response = session.send(request.prepare())
            response.raise_for_status()
            return response.json()["data"]
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                raise UnauthorizedError(e)
            raise

    def new(
        self,
        model_name: str,
        model_params: Dict[str, Any],
        model_info: Optional[Dict[str, Any]] = None,
        return_request: bool = False,
    ) -> Union[Dict[str, Any], requests.Request]:
        """
        Add a new model to the proxy.

        Args:
            model_name (str): Name of the model to add
            model_params (Dict[str, Any]): Parameters for the model (e.g., model type, api_base, api_key)
            model_info (Optional[Dict[str, Any]]): Additional information about the model
            return_request (bool): If True, returns the prepared request object instead of executing it

        Returns:
            Union[Dict[str, Any], requests.Request]: Either the response from the server or
            a prepared request object if return_request is True

        Raises:
            UnauthorizedError: If the request fails with a 401 status code
            requests.exceptions.RequestException: If the request fails with any other error
        """
        url = f"{self._base_url}/model/new"

        data = {
            "model_name": model_name,
            "litellm_params": model_params,
        }
        if model_info:
            data["model_info"] = model_info

        request = requests.Request("POST", url, headers=self._get_headers(), json=data)

        if return_request:
            return request

        # Prepare and send the request
        session = requests.Session()
        try:
            response = session.send(request.prepare())
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                raise UnauthorizedError(e)
            raise

    def delete(self, model_id: str, return_request: bool = False) -> Union[Dict[str, Any], requests.Request]:
        """
        Delete a model from the proxy.

        Args:
            model_id (str): ID of the model to delete (e.g., "2f23364f-4579-4d79-a43a-2d48dd551c2e")
            return_request (bool): If True, returns the prepared request object instead of executing it

        Returns:
            Union[Dict[str, Any], requests.Request]: Either the response from the server or
            a prepared request object if return_request is True

        Raises:
            UnauthorizedError: If the request fails with a 401 status code
            NotFoundError: If the request fails with a 404 status code or indicates the model was not found
            requests.exceptions.RequestException: If the request fails with any other error
        """
        url = f"{self._base_url}/model/delete"
        data = {"id": model_id}

        request = requests.Request("POST", url, headers=self._get_headers(), json=data)

        if return_request:
            return request

        # Prepare and send the request
        session = requests.Session()
        try:
            response = session.send(request.prepare())
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                raise UnauthorizedError(e)
            if e.response.status_code == 404 or "not found" in e.response.text.lower():
                raise NotFoundError(e)
            raise

    def get(
        self, model_id: Optional[str] = None, model_name: Optional[str] = None, return_request: bool = False
    ) -> Union[Dict[str, Any], requests.Request]:
        """
        Get information about a specific model by its ID or name.

        Args:
            model_id (Optional[str]): ID of the model to retrieve
            model_name (Optional[str]): Name of the model to retrieve
            return_request (bool): If True, returns the prepared request object instead of executing it

        Returns:
            Union[Dict[str, Any], requests.Request]: Either the model information from the server or
            a prepared request object if return_request is True

        Raises:
            ValueError: If neither model_id nor model_name is provided, or if both are provided
            UnauthorizedError: If the request fails with a 401 status code
            NotFoundError: If the model is not found
            requests.exceptions.RequestException: If the request fails with any other error
        """
        if (model_id is None and model_name is None) or (model_id is not None and model_name is not None):
            raise ValueError("Exactly one of model_id or model_name must be provided")

        # If return_request is True, delegate to info
        if return_request:
            result = self.info(return_request=True)
            assert isinstance(result, requests.Request)
            return result

        # Get all models and filter
        models = self.info()
        assert isinstance(models, List)

        # Find the matching model
        for model in models:
            if (model_id and model.get("model_info", {}).get("id") == model_id) or (
                model_name and model.get("model_name") == model_name
            ):
                return model

        # If we get here, no model was found
        if model_id:
            msg = f"Model with id={model_id} not found"
        elif model_name:
            msg = f"Model with model_name={model_name} not found"
        else:
            msg = "Unknown error trying to find model"
        raise NotFoundError(
            requests.exceptions.HTTPError(
                msg,
                response=requests.Response(),  # Empty response since we didn't make a direct request
            )
        )

    def info(self, return_request: bool = False) -> Union[List[Dict[str, Any]], requests.Request]:
        """
        Get detailed information about all models from the server.

        Args:
            return_request (bool): If True, returns the prepared request object instead of executing it

        Returns:
            Union[List[Dict[str, Any]], requests.Request]: Either a list of model information dictionaries
            or a prepared request object if return_request is True

        Raises:
            UnauthorizedError: If the request fails with a 401 status code
            requests.exceptions.RequestException: If the request fails with any other error
        """
        url = f"{self._base_url}/v1/model/info"
        request = requests.Request("GET", url, headers=self._get_headers())

        if return_request:
            return request

        # Prepare and send the request
        session = requests.Session()
        try:
            response = session.send(request.prepare())
            response.raise_for_status()
            return response.json()["data"]
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                raise UnauthorizedError(e)
            raise

    def update(
        self,
        model_id: str,
        model_params: Dict[str, Any],
        model_info: Optional[Dict[str, Any]] = None,
        return_request: bool = False,
    ) -> Union[Dict[str, Any], requests.Request]:
        """
        Update an existing model's configuration.

        Args:
            model_id (str): ID of the model to update
            model_params (Dict[str, Any]): New parameters for the model (e.g., model type, api_base, api_key)
            model_info (Optional[Dict[str, Any]]): Additional information about the model
            return_request (bool): If True, returns the prepared request object instead of executing it

        Returns:
            Union[Dict[str, Any], requests.Request]: Either the response from the server or
            a prepared request object if return_request is True

        Raises:
            UnauthorizedError: If the request fails with a 401 status code
            NotFoundError: If the model is not found
            requests.exceptions.RequestException: If the request fails with any other error
        """
        url = f"{self._base_url}/model/update"

        data = {
            "id": model_id,
            "litellm_params": model_params,
        }
        if model_info:
            data["model_info"] = model_info

        request = requests.Request("POST", url, headers=self._get_headers(), json=data)

        if return_request:
            return request

        # Prepare and send the request
        session = requests.Session()
        try:
            response = session.send(request.prepare())
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                raise UnauthorizedError(e)
            if e.response.status_code == 404 or "not found" in e.response.text.lower():
                raise NotFoundError(e)
            raise