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

import requests
from pydantic_v1 import Extra, root_validator

from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM
from langchain.llms.utils import enforce_stop_tokens
from langchain.utils import get_from_dict_or_env

DEFAULT_NEBULA_SERVICE_URL = "https://api-nebula.symbl.ai"
DEFAULT_NEBULA_SERVICE_PATH = "/v1/model/generate"

logger = logging.getLogger(__name__)


class Nebula(LLM):
    """Nebula Service models.

    To use, you should have the environment variable ``NEBULA_SERVICE_URL``,
    ``NEBULA_SERVICE_PATH`` and ``NEBULA_SERVICE_API_KEY`` set with your Nebula
    Service, or pass it as a named parameter to the constructor.

    Example:
        .. code-block:: python

            from langchain.llms import Nebula

            nebula = Nebula(
                nebula_service_url="SERVICE_URL",
                nebula_service_path="SERVICE_ROUTE",
                nebula_api_key="SERVICE_TOKEN",
            )
    """  # noqa: E501

    """Key/value arguments to pass to the model. Reserved for future use"""
    model_kwargs: Optional[dict] = None

    """Optional"""
    nebula_service_url: Optional[str] = None
    nebula_service_path: Optional[str] = None
    nebula_api_key: Optional[str] = None
    conversation: str = ""
    return_scores: Optional[str] = "false"
    max_new_tokens: Optional[int] = 2048
    top_k: Optional[float] = 2
    penalty_alpha: Optional[float] = 0.1

    class Config:
        """Configuration for this pydantic object."""

        extra = Extra.forbid

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate that api key and python package exists in environment."""
        nebula_service_url = get_from_dict_or_env(
            values,
            "nebula_service_url",
            "NEBULA_SERVICE_URL",
            DEFAULT_NEBULA_SERVICE_URL,
        )
        nebula_service_path = get_from_dict_or_env(
            values,
            "nebula_service_path",
            "NEBULA_SERVICE_PATH",
            DEFAULT_NEBULA_SERVICE_PATH,
        )
        nebula_api_key = get_from_dict_or_env(
            values, "nebula_api_key", "NEBULA_SERVICE_API_KEY", ""
        )

        if nebula_service_url.endswith("/"):
            nebula_service_url = nebula_service_url[:-1]
        if not nebula_service_path.startswith("/"):
            nebula_service_path = "/" + nebula_service_path

        """ TODO: Future login"""
        """
        try:
            nebula_service_endpoint = f"{nebula_service_url}{nebula_service_path}"
            headers = {
                "Content-Type": "application/json",
                "ApiKey": "{nebula_api_key}",
                }
            requests.get(nebula_service_endpoint, headers=headers)
        except requests.exceptions.RequestException as e:
            raise ValueError(e)
        """

        values["nebula_service_url"] = nebula_service_url
        values["nebula_service_path"] = nebula_service_path
        values["nebula_api_key"] = nebula_api_key

        return values

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        _model_kwargs = self.model_kwargs or {}
        return {
            "nebula_service_url": self.nebula_service_url,
            "nebula_service_path": self.nebula_service_path,
            **{"model_kwargs": _model_kwargs},
            "conversation": self.conversation,
        }

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

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        """Call out to Nebula Service endpoint.
        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 = nebula("Tell me a joke.")
        """

        _model_kwargs = self.model_kwargs or {}

        nebula_service_endpoint = f"{self.nebula_service_url}{self.nebula_service_path}"

        headers = {
            "Content-Type": "application/json",
            "ApiKey": f"{self.nebula_api_key}",
        }

        body = {
            "prompt": {
                "instruction": prompt,
                "conversation": {"text": f"{self.conversation}"},
            },
            "return_scores": self.return_scores,
            "max_new_tokens": self.max_new_tokens,
            "top_k": self.top_k,
            "penalty_alpha": self.penalty_alpha,
        }

        if len(self.conversation) == 0:
            raise ValueError("Error conversation is empty.")

        logger.debug(f"NEBULA _model_kwargs: {_model_kwargs}")
        logger.debug(f"NEBULA body: {body}")
        logger.debug(f"NEBULA kwargs: {kwargs}")
        logger.debug(f"NEBULA conversation: {self.conversation}")

        # call API
        try:
            response = requests.post(
                nebula_service_endpoint, headers=headers, json=body
            )
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Error raised by inference endpoint: {e}")

        logger.debug(f"NEBULA response: {response}")

        if response.status_code != 200:
            raise ValueError(
                f"Error returned by service, status code {response.status_code}"
            )

        """ get the result """
        text = response.text

        """ enforce stop """
        if stop is not None:
            # This is required since the stop tokens
            # are not enforced by the model parameters
            text = enforce_stop_tokens(text, stop)

        return text
