import logging
from typing import Any, Dict, List, Optional, Iterator
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM

from langchain.schema.output import GenerationChunk
from langchain.pydantic_v1 import Field, root_validator
from langchain.utils import (
    get_from_dict_or_env,
    get_pydantic_field_names,
)

from turing_planet.api.spark_chat_client import SparkLLMClient

logger = logging.getLogger(__name__)


def _convert_prompt_to_dict(
        prompt: str,
) -> dict:
    return {"role": "user", "content": prompt}


class SparkAI(LLM):
    client: Any = None  #: :meta private:
    endpoint: Optional[str] = None

    trace_id: str = "langchain_user"
    max_tokens: int = 4096
    temperature: float = 0.5
    top_k: int = 4
    adjust_tokens: bool = True
    domain: str = ""

    streaming: bool = False
    request_timeout: int = 120

    model_kwargs: Dict[str, Any] = Field(default_factory=dict)

    @classmethod
    def is_lc_serializable(cls) -> bool:
        """Return whether this model can be serialized by Langchain."""
        return True

    @root_validator(pre=True)
    def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        """Build extra kwargs from additional params that were passed in."""
        all_required_field_names = get_pydantic_field_names(cls)
        extra = values.get("model_kwargs", {})
        for field_name in list(values):
            if field_name in extra:
                raise ValueError(f"Found {field_name} supplied twice.")
            if field_name not in all_required_field_names:
                logger.warning(
                    f"""WARNING! {field_name} is not default parameter.
                    {field_name} was transferred to model_kwargs.
                    Please confirm that {field_name} is what you intended."""
                )
                extra[field_name] = values.pop(field_name)

        invalid_model_kwargs = all_required_field_names.intersection(extra.keys())
        if invalid_model_kwargs:
            raise ValueError(
                f"Parameters {invalid_model_kwargs} should be specified explicitly. "
                f"Instead they were passed in as part of `model_kwargs` parameter."
            )

        values["model_kwargs"] = extra

        return values

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        values["endpoint"] = get_from_dict_or_env(
            values,
            "endpoint",
            "TURING_PLANET_ENDPOINT",
            "127.0.0.1:9980",
        )
        # put extra params into model_kwargs
        values["model_kwargs"]["max_tokens"] = values["max_tokens"]
        values["model_kwargs"]["temperature"] = values["temperature"]
        values["model_kwargs"]["top_k"] = values["top_k"]
        values["model_kwargs"]["adjust_tokens"] = values["adjust_tokens"]
        values["model_kwargs"]["domain"] = values["domain"]

        values["client"] = SparkLLMClient(
            endpoint=values["endpoint"],
            trace_id=values["trace_id"],
            model_kwargs=values["model_kwargs"],
        )
        return values

    def _call(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> str:
        """Call out to an spark models endpoint for each generation with a prompt.
        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 = SparkAI("Tell me a joke.")
        """
        if self.streaming:
            completion = ""
            for chunk in self._stream(prompt, stop, run_manager, **kwargs):
                completion += chunk.text
            return completion

        self.client.arun(
            [_convert_prompt_to_dict(prompt)],
            self.model_kwargs,
            kwargs,
            False,
        )
        completion = {}
        for content in self.client.subscribe(timeout=self.request_timeout):
            if "data" not in content:
                continue
            completion = content["data"]
        return completion["content"]

    def _stream(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> Iterator[GenerationChunk]:

        self.client.arun(
            [_convert_prompt_to_dict(prompt)],
            self.model_kwargs,
            kwargs,
            True,
        )

        for content in self.client.subscribe(timeout=self.request_timeout):
            if "data" not in content:
                continue
            delta = content["data"]
            chunk = GenerationChunk(text=delta["content"])
            yield chunk
            if run_manager:
                run_manager.on_llm_new_token(chunk.text)

    @property
    def _llm_type(self) -> str:
        return "spark-llm"
