from typing import Any, Dict, List, Optional

import requests
from pydantic import BaseModel, Extra, root_validator

from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM

class Guanaco(LLM):
    """Wrapper around AI21 large language models.

    To use, you should have the environment variable ``AI21_API_KEY``
    set with your API key.

    Example:
        .. code-block:: python
    """
    model: str = "baichuan"
    """Model name to use."""

    temperature: float = 0.7
    """What sampling temperature to use."""

    maxTokens: int = 256
    """The maximum number of tokens to generate in the completion."""

    minTokens: int = 0
    """The minimum number of tokens to generate in the completion."""

    topP: float = 1.0
    """Total probability mass of tokens to consider at each step."""

    base_url: Optional[str] = 'http://ecsb.crcloud.com/ecsb/gw/sys/rf?ssdp=QXBpX0lEPWNybGFuZC5jcmxhcnNzLlNDUjAxLjAwMDEmQXBpX1ZlcnNpb249MS4wJkFwcF9TdWJfSUQ9MDAwNTAwMUkwMDAyJkFwcF9Ub2tlbj04M2YwMDE3OTY3NjE0MGRmOWZiNDA0Y2MwNGM3YWRkNyZQYXJ0bmVyX0lEPTAwMDUwMDAwJlNpZ249MjA5MzQzODJkYThkNGJjNGE2OWJhZmQ2MGQxYTJmNzUmU3lzX0lEPTAwMDUwMDFJJlRpbWVfU3RhbXA9MjAyMy0wNS0yMiAxNjowMjoxMjowNjgmVXNlcl9Ub2tlbj0='

    """Base url to use, if None decides based on model name."""

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

        extra = Extra.forbid


    @property
    def _default_params(self) -> Dict[str, Any]:
        """Get the default parameters for calling AI21 API."""
        return {
            "temperature": self.temperature,
            "maxTokens": self.maxTokens,
            "minTokens": self.minTokens,
            "topP": self.topP,
        }

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """Get the identifying parameters."""
        return {**{"model": self.model}, **self._default_params}

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

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
    ) -> str:
        """Call out to AI21's complete 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.
        """

        response = requests.post(
            url= self.base_url,
            json={
                "from_server": "646c52e724c3203f886e4bf8a22eddb499e03972", 
                "model": self.model, 
                "message": prompt, 
                "raw_text": "1"
            }
        )
        if response.status_code != 200:
            optional_detail = response.json().get("error")
            raise ValueError(
                f"complete call failed with status code {response.status_code}."
                f" Details: {optional_detail}"
            )
        try:
            response_json = response.json()
            #print(response_json)
            return response_json["data"]
        except Exception as e:
            return e
