from typing import Any, Dict, List, Mapping, Optional

import requests
from pydantic import root_validator

try:
    from langchain.callbacks.manager import CallbackManagerForLLMRun
    from langchain.llms.base import LLM
    from langchain.llms.utils import enforce_stop_tokens
except Exception as ex:
    raise Exception(
        "Failed to import langchain."
        " Please install langchain by using `pip install langchain` command"
    ) from ex


class TruefoundryLLM(LLM):
    """Wrapper around TFY model deployment.

    To use this class, you need to have the langchain library installed.

    Example:
        .. code-block:: python

            from servicefoundry.langchain import TruefoundryLLM
            endpoint_url = (
                "https://pythia-70m-model-model-catalogue.demo2.truefoundry.tech"
            )
            model = TruefoundryLLM(
                endpoint_url=endpoint_url,
                parameters={
                    "max_new_tokens": 100,
                    "temperature": 0.7,
                    "top_k": 5,
                    "top_p": 0.9
                }
            )
    """

    endpoint_url: str
    model_name: Optional[str] = None
    parameters: Optional[dict] = None

    @root_validator(pre=False)
    def validate_model_name(cls, values: Dict):
        endpoint_url = values["endpoint_url"]
        model_name = values.get("model_name")

        url = f"{endpoint_url.strip('/')}/v2/repository/index"
        try:
            response = requests.post(url, json={})
            response.raise_for_status()
        except Exception as ex:
            raise Exception(f"Error raised by inference API: {ex}") from ex
        models = response.json()

        if len(models) == 0:
            raise ValueError("No model is deployed in the model server")

        model_names = [m.get("name") for m in models]

        if model_name and model_name not in model_names:
            raise ValueError(
                f"Model {model_name!r} is not available in the model server. Available models {model_names!r}"
            )

        if not model_name and len(model_names) > 1:
            raise ValueError(
                f"Please pass `model_name` while instantiating `TruefoundryLLM`. Available models are {model_names!r} "
            )

        if model_name:
            return values

        values["model_name"] = model_names[0]
        return values

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {
            "endpoint_url": self.endpoint_url,
            "model_name": self.model_name,
        }

    @property
    def _llm_type(self) -> str:
        """Return type of llm."""
        return "tfy_model_deployment"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **params: Any,
    ) -> str:
        """Call out to the deployed model

        Args:
            prompt: The prompt to pass into the model.
            stop: Optional list of stop words to use when generating.

        Returns:
            The string generated by the model.

        Example:
            .. code-block:: python

                response = model("Tell me a joke.")
        """
        _params_already_set = self.parameters or {}

        params = {**_params_already_set, **params}
        params["return_full_text"] = False

        payload = {"inputs": prompt, "parameters": params}
        url = f"{self.endpoint_url.strip('/')}/v2/models/{self.model_name}/infer/simple"

        try:
            response = requests.post(url, json=payload)
            response.raise_for_status()
        except Exception as ex:
            raise Exception(f"Error raised by inference API: {ex}") from ex
        response_dict = response.json()
        if "error" in response_dict:
            raise ValueError(
                f"Error raised by inference API: {response_dict['error']!r}"
            )

        inference_result = response_dict[0]
        text = ""

        if "generated_text" in inference_result:
            text = inference_result["generated_text"]
        elif "summarization" in inference_result:
            text = inference_result["summary_text"]
        else:
            raise ValueError(f"Could not parse inference response: {response_dict!r}")

        if stop is not None:
            text = enforce_stop_tokens(text, stop)

        return text
