from __future__ import annotations

import json
import warnings
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional

from langchain_core.callbacks.manager import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain_core.language_models.llms import LLM
from langchain_core.outputs import GenerationChunk
from pydantic import ConfigDict, Field, PrivateAttr

from langchain_nvidia_ai_endpoints._common import _NVIDIAClient
from langchain_nvidia_ai_endpoints._statics import Model

_DEFAULT_MODEL_NAME: str = "nvidia/mistral-nemo-minitron-8b-base"


class NVIDIA(LLM):
    """LangChain LLM that uses the Completions API with NVIDIA NIMs."""

    model_config = ConfigDict(
        validate_assignment=True,
    )

    _client: _NVIDIAClient = PrivateAttr()

    _default_model_name: str = "nvidia/mistral-nemo-minitron-8b-base"

    base_url: Optional[str] = Field(
        default=None,
        description="Base url for model listing and invocation",
    )

    model: Optional[str] = Field(None, description="The model to use for completions.")

    _init_args: Dict[str, Any] = PrivateAttr()
    """Stashed arguments given to the constructor that can be passed to
    the Completions API endpoint."""

    def __check_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
        """Check kwargs, warn for unknown keys, and return a copy recognized keys."""
        completions_arguments = {
            "frequency_penalty",
            "max_tokens",
            "presence_penalty",
            "seed",
            "stop",
            "temperature",
            "top_p",
            "best_of",
            "echo",
            "logit_bias",
            "logprobs",
            "n",
            "suffix",
            "user",
            "stream",
        }

        recognized_kwargs = {
            k: v for k, v in kwargs.items() if k in completions_arguments
        }
        unrecognized_kwargs = set(kwargs) - completions_arguments
        if len(unrecognized_kwargs) > 0:
            warnings.warn(f"Unrecognized, ignored arguments: {unrecognized_kwargs}")

        return recognized_kwargs

    def __init__(
        self,
        *,
        nvidia_api_key: Optional[str] = None,
        api_key: Optional[str] = None,
        **kwargs: Any,
    ):
        """Create a new NVIDIA LLM for Completions APIs.

        This class provides access to a NVIDIA NIM for completions. By default, it
        connects to a hosted NIM, but can be configured to connect to a local NIM
        using the `base_url` parameter.

        An API key is required to connect to the hosted NIM.

        Args:
            nvidia_api_key: The API key to use for connecting to the hosted NIM.
            api_key: Alternative to `nvidia_api_key`.
            **kwargs: Additional parameters passed to the underlying client.

        The recommended way to provide the API key is through the `NVIDIA_API_KEY`
        environment variable.

        Additional arguments that can be passed to the Completions API:

        - max_tokens (int): The maximum number of tokens to generate.
        - stop (str or List[str]): The stop sequence to use for generating completions.
        - temperature (float): The temperature to use for generating completions.
        - top_p (float): The top-p value to use for generating completions.
        - frequency_penalty (float): The frequency penalty to apply to the completion.
        - presence_penalty (float): The presence penalty to apply to the completion.
        - seed (int): The seed to use for generating completions.

        These additional arguments can also be passed with `bind()`, e.g.
        `NVIDIA().bind(max_tokens=512)`, or pass directly to `invoke()` or `stream()`,
        e.g. `NVIDIA().invoke("prompt", max_tokens=512)`.
        """
        super().__init__(**kwargs)

        # allow nvidia_base_url as an alternative for base_url
        base_url = kwargs.pop("nvidia_base_url", self.base_url)

        # allow nvidia_api_key as an alternative for api_key
        api_key = nvidia_api_key or api_key

        # Extract verify_ssl from kwargs, default to True
        verify_ssl = kwargs.pop("verify_ssl", True)

        self._client = _NVIDIAClient(
            **({"base_url": base_url} if base_url else {}),  # only pass if set
            mdl_name=self.model,
            default_hosted_model_name=_DEFAULT_MODEL_NAME,
            **({"api_key": api_key} if api_key else {}),  # only pass if set
            infer_path="{base_url}/completions",
            cls=self.__class__.__name__,
            verify_ssl=verify_ssl,
        )
        # todo: only store the model in one place
        # the model may be updated to a newer name during initialization
        self.model = self._client.mdl_name
        # same for base_url
        self.base_url = self._client.base_url

        # stash all additional args that can be passed to the Completions API,
        # but first make sure we pull out any args that are processed elsewhere.
        for key in [
            "model",
            "nvidia_base_url",
            "base_url",
        ]:
            if key in kwargs:
                del kwargs[key]
        self._init_args = self.__check_kwargs(kwargs)

    @property
    def available_models(self) -> List[Model]:
        """Get a list of available models that work with NVIDIA."""
        return self._client.get_available_models(self.__class__.__name__)

    @classmethod
    def get_available_models(
        cls,
        **kwargs: Any,
    ) -> List[Model]:
        """Get a list of available models that work with the Completions API."""
        return cls(**kwargs).available_models

    @property
    def _llm_type(self) -> str:
        """Get the type of language model used by this chat model.

        Used for logging purposes only.
        """
        return "NVIDIA"

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """Get parameters used to help identify the LLM."""
        return {
            "model": self.model,
            "base_url": self.base_url,
        }

    def _prepare_call_payload(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """Prepare payload for non-streaming calls (both sync and async).

        Args:
            prompt: The prompt to send
            stop: Stop words
            kwargs: Additional keyword arguments

        Returns:
            Payload dictionary
        """
        payload: Dict[str, Any] = {
            "model": self.model,
            "prompt": prompt,
            **self._init_args,
            **self.__check_kwargs(kwargs),
        }
        if stop:
            payload["stop"] = stop

        if payload.get("stream", False):
            warnings.warn("stream set to true for non-streaming call, ignoring")
            del payload["stream"]

        return payload

    def _prepare_stream_payload(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """Prepare payload for streaming calls (both sync and async).

        Args:
            prompt: The prompt to send
            stop: Stop words
            kwargs: Additional keyword arguments

        Returns:
            Payload dictionary
        """
        payload: Dict[str, Any] = {
            "model": self.model,
            "prompt": prompt,
            "stream": True,
            **self._init_args,
            **self.__check_kwargs(kwargs),
        }
        if stop:
            payload["stop"] = stop

        # we construct payload w/ **kwargs positioned to override stream=True,
        # this lets us know if a user passed stream=False
        if not payload.get("stream", True):
            warnings.warn("stream set to false for streaming call, ignoring")
            payload["stream"] = True

        return payload

    def _process_result(self, result: Dict[str, Any]) -> str:
        """Process parsed JSON result from both sync and async call methods.

        Args:
            result: Parsed JSON response

        Returns:
            Generated text
        """
        # todo: handle response's usage and system_fingerprint
        choices = result["choices"]
        # todo: write a test for this by setting n > 1 on the request
        #       aug 2024: n > 1 is not supported by endpoints
        if len(choices) > 1:
            warnings.warn(
                f"Multiple choices in response, returning only the first: {choices}"
            )

        return choices[0]["text"]

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        payload = self._prepare_call_payload(prompt, stop, **kwargs)
        response = self._client.get_req(payload=payload)
        response.raise_for_status()
        result = response.json()
        return self._process_result(result)

    def _stream(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[GenerationChunk]:
        payload = self._prepare_stream_payload(prompt, stop, **kwargs)
        for chunk in self._client.get_req_stream(payload=payload):
            content = chunk["content"]
            generation = GenerationChunk(text=content)
            if run_manager:  # todo: add tests for run_manager
                run_manager.on_llm_new_token(content, chunk=generation)
            yield generation

    async def _acall(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        payload = self._prepare_call_payload(prompt, stop, **kwargs)
        response_text = await self._client.aget_req(payload=payload)
        result = json.loads(response_text)
        return self._process_result(result)

    async def _astream(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> AsyncIterator[GenerationChunk]:
        payload = self._prepare_stream_payload(prompt, stop, **kwargs)
        async for chunk in self._client.aget_req_stream(payload=payload):
            content = chunk["content"]
            generation = GenerationChunk(text=content)
            if run_manager:
                await run_manager.on_llm_new_token(content, chunk=generation)
            yield generation
