"""Wrapper around Cerebras' Chat Completions API."""

from typing import Any, Dict, List, Optional

import openai
from langchain_core.language_models.chat_models import LangSmithParams
from langchain_core.utils import (
    from_env,
    secret_from_env,
)

# We ignore the "unused imports" here since we want to reexport these from this package.
from langchain_openai.chat_models.base import (
    BaseChatOpenAI,
)
from pydantic import Field, SecretStr, model_validator
from typing_extensions import Self

CEREBRAS_BASE_URL = "https://api.cerebras.ai/v1/"


class ChatCerebras(BaseChatOpenAI):
    r"""ChatCerebras chat model.

    Setup:
        Install ``langchain-cerebras`` and set environment variable ``CEREBRAS_API_KEY``.

        .. code-block:: bash

            pip install -U langchain-cerebras
            export CEREBRAS_API_KEY="your-api-key"


    Key init args — completion params:
        model: str
            Name of model to use.
        temperature: Optional[float]
            Sampling temperature.
        max_tokens: Optional[int]
            Max number of tokens to generate.

    Key init args — client params:
        timeout: Union[float, Tuple[float, float], Any, None]
            Timeout for requests.
        max_retries: Optional[int]
            Max number of retries.
        api_key: Optional[str]
            Cerebras API key. If not passed in will be read from env var CEREBRAS_API_KEY.

    Instantiate:
        .. code-block:: python

            from langchain_cerebras import ChatCerebras

            llm = ChatCerebras(
                model="llama-3.3-70b",
                temperature=0,
                max_tokens=None,
                timeout=None,
                max_retries=2,
                # api_key="...",
                # other params...
            )

    Invoke:
        .. code-block:: python

            messages = [
                (
                    "system",
                    "You are a helpful translator. Translate the user sentence to French.",
                ),
                ("human", "I love programming."),
            ]
            llm.invoke(messages)

        .. code-block:: python

            AIMessage(
                content='The translation of "I love programming" to French is:\n\n"J\'adore programmer."',
                response_metadata={
                    'token_usage': {'completion_tokens': 20, 'prompt_tokens': 32, 'total_tokens': 52},
                    'model_name': 'llama-3.3-70b',
                    'system_fingerprint': 'fp_679dff74c0',
                    'finish_reason': 'stop',
                },
                id='run-377c2887-30ef-417e-b0f5-83efc8844f12-0',
                usage_metadata={'input_tokens': 32, 'output_tokens': 20, 'total_tokens': 52})

    Stream:
        .. code-block:: python

            for chunk in llm.stream(messages):
                print(chunk)

        .. code-block:: python

            content='' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content='The' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=' translation' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=' of' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=' "' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content='I' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=' love' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=' programming' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content='"' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=' to' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=' French' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=' is' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=':\n\n' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content='"' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content='J' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content="'" id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content='ad' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content='ore' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content=' programmer' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content='."' id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'
            content='' response_metadata={'finish_reason': 'stop', 'model_name': 'llama-3.3-70b', 'system_fingerprint': 'fp_679dff74c0'} id='run-3f9dc84e-208f-48da-b15d-e552b6759c24'

    Async:
        .. code-block:: python

            await llm.ainvoke(messages)

            # stream:
            # async for chunk in (await llm.astream(messages))

            # batch:
            # await llm.abatch([messages])

        .. code-block:: python

            AIMessage(
                content='The translation of "I love programming" to French is:\n\n"J\'adore programmer."',
                response_metadata={
                    'token_usage': {'completion_tokens': 20, 'prompt_tokens': 32, 'total_tokens': 52},
                    'model_name': 'llama-3.3-70b',
                    'system_fingerprint': 'fp_679dff74c0',
                    'finish_reason': 'stop',
                },
                id='run-377c2887-30ef-417e-b0f5-83efc8844f12-0',
                usage_metadata={'input_tokens': 32, 'output_tokens': 20, 'total_tokens': 52})

    Tool calling:
        .. code-block:: python

            from langchain_core.pydantic_v1 import BaseModel, Field

            llm = ChatCerebras(model="llama-3.3-70b")

            class GetWeather(BaseModel):
                '''Get the current weather in a given location'''

                location: str = Field(
                    ..., description="The city and state, e.g. San Francisco, CA"
                )

            class GetPopulation(BaseModel):
                '''Get the current population in a given location'''

                location: str = Field(
                    ..., description="The city and state, e.g. San Francisco, CA"
                )

            llm_with_tools = llm.bind_tools([GetWeather, GetPopulation])
            ai_msg = llm_with_tools.invoke(
                "Which city is bigger: LA or NY?"
            )
            ai_msg.tool_calls


        .. code-block:: python

            [
                {
                    'name': 'GetPopulation',
                    'args': {'location': 'NY'},
                    'id': 'call_m5tstyn2004pre9bfuxvom8x',
                    'type': 'tool_call'
                },
                {
                    'name': 'GetPopulation',
                    'args': {'location': 'LA'},
                    'id': 'call_0vjgq455gq1av5sp9eb1pw6a',
                    'type': 'tool_call'
                }
            ]

    Structured output:
        .. code-block:: python

            from typing import Optional

            from langchain_core.pydantic_v1 import BaseModel, Field


            class Joke(BaseModel):
                '''Joke to tell user.'''

                setup: str = Field(description="The setup of the joke")
                punchline: str = Field(description="The punchline to the joke")
                rating: Optional[int] = Field(description="How funny the joke is, from 1 to 10")


            structured_llm = llm.with_structured_output(Joke)
            structured_llm.invoke("Tell me a joke about cats")

        .. code-block:: python

            Joke(
                setup='Why was the cat sitting on the computer?',
                punchline='To keep an eye on the mouse!',
                rating=7
            )

    JSON mode:
        .. code-block:: python

            json_llm = llm.bind(response_format={"type": "json_object"})
            ai_msg = json_llm.invoke(
                "Return a JSON object with key 'random_ints' and a value of 10 random ints in [0-99]"
            )
            ai_msg.content

        .. code-block:: python

            ' {\\n"random_ints": [\\n13,\\n54,\\n78,\\n45,\\n67,\\n90,\\n11,\\n29,\\n84,\\n33\\n]\\n}'

    Token usage:
        .. code-block:: python

            ai_msg = llm.invoke(messages)
            ai_msg.usage_metadata

        .. code-block:: python

            {'input_tokens': 37, 'output_tokens': 6, 'total_tokens': 43}

    Response metadata
        .. code-block:: python

            ai_msg = llm.invoke(messages)
            ai_msg.response_metadata

        .. code-block:: python

            {
                'token_usage': {
                    'completion_tokens': 4,
                    'prompt_tokens': 19,
                    'total_tokens': 23
                    },
                'model_name': 'mistralai/Mixtral-8x7B-Instruct-v0.1',
                'system_fingerprint': None,
                'finish_reason': 'eos',
                'logprobs': None
            }

    """  # noqa: E501

    @property
    def lc_secrets(self) -> Dict[str, str]:
        """A map of constructor argument names to secret ids.

        For example,
            {"cerebras_api_key": "CEREBRAS_API_KEY"}
        """
        return {"cerebras_api_key": "CEREBRAS_API_KEY"}

    @classmethod
    def get_lc_namespace(cls) -> List[str]:
        """Get the namespace of the langchain object."""
        return ["langchain", "chat_models", "cerebras"]

    @property
    def lc_attributes(self) -> Dict[str, Any]:
        """List of attribute names that should be included in the serialized kwargs.

        These attributes must be accepted by the constructor.
        """
        attributes: Dict[str, Any] = {}

        if self.cerebras_api_base:
            attributes["cerebras_api_base"] = self.cerebras_api_base

        if self.cerebras_proxy:
            attributes["cerebras_proxy"] = self.cerebras_proxy

        return attributes

    @property
    def _llm_type(self) -> str:
        """Return type of chat model."""
        return "cerebras-chat"

    def _get_ls_params(
        self, stop: Optional[List[str]] = None, **kwargs: Any
    ) -> LangSmithParams:
        """Get the parameters used to invoke the model."""
        params = super()._get_ls_params(stop=stop, **kwargs)
        params["ls_provider"] = "cerebras"
        return params

    model_name: str = Field(alias="model")
    """Model name to use."""
    cerebras_api_key: Optional[SecretStr] = Field(
        alias="api_key",
        default_factory=secret_from_env("CEREBRAS_API_KEY", default=None),
    )
    """Automatically inferred from env are `CEREBRAS_API_KEY` if not provided."""
    cerebras_api_base: str = Field(
        default_factory=from_env("CEREBRAS_API_BASE", default=CEREBRAS_BASE_URL),
        alias="base_url",
    )

    cerebras_proxy: str = Field(default_factory=from_env("CEREBRAS_PROXY", default=""))

    @model_validator(mode="after")
    def validate_environment(self) -> Self:
        """Validate that api key and python package exists in environment."""
        if self.n is not None and self.n < 1:
            raise ValueError("n must be at least 1.")
        if self.n is not None and self.n > 1 and self.streaming:
            raise ValueError("n must be 1 when streaming.")

        client_params = {
            "api_key": (
                self.cerebras_api_key.get_secret_value()
                if self.cerebras_api_key
                else None
            ),
            # Ensure we always fallback to the Cerebras API url.
            "base_url": self.cerebras_api_base,
            "timeout": self.request_timeout,
            "default_headers": self.default_headers,
            "default_query": self.default_query,
        }
        if self.max_retries is not None:
            client_params["max_retries"] = self.max_retries

        if self.cerebras_proxy and (self.http_client or self.http_async_client):
            raise ValueError(
                "Cannot specify 'cerebras_proxy' if one of "
                "'http_client'/'http_async_client' is already specified. Received:\n"
                f"{self.cerebras_proxy=}\n{self.http_client=}\n{self.http_async_client=}"
            )
        if not self.client:
            if self.cerebras_proxy and not self.http_client:
                try:
                    import httpx
                except ImportError as e:
                    raise ImportError(
                        "Could not import httpx python package. "
                        "Please install it with `pip install httpx`."
                    ) from e
                self.http_client = httpx.Client(proxy=self.cerebras_proxy)
            sync_specific = {"http_client": self.http_client}
            self.root_client = openai.OpenAI(**client_params, **sync_specific)  # type: ignore
            self.client = self.root_client.chat.completions
        if not self.async_client:
            if self.cerebras_proxy and not self.http_async_client:
                try:
                    import httpx
                except ImportError as e:
                    raise ImportError(
                        "Could not import httpx python package. "
                        "Please install it with `pip install httpx`."
                    ) from e
                self.http_async_client = httpx.AsyncClient(proxy=self.cerebras_proxy)
            async_specific = {"http_client": self.http_async_client}
            self.root_async_client = openai.AsyncOpenAI(
                **client_params,  # type: ignore
                **async_specific,  # type: ignore
            )
            self.async_client = self.root_async_client.chat.completions
        return self