"""SiliconFlow chat models."""

from __future__ import annotations

import json
from collections.abc import Iterator
from json import JSONDecodeError
from typing import Any, Optional, TypeVar, Union

import openai
from langchain_core.callbacks import (
    CallbackManagerForLLMRun,
)
from langchain_core.language_models import LangSmithParams, LanguageModelInput
from langchain_core.messages import AIMessageChunk, BaseMessage
from langchain_core.outputs import ChatGenerationChunk, ChatResult
from langchain_core.utils import from_env, secret_from_env
from langchain_openai.chat_models.base import BaseChatOpenAI
from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
from typing_extensions import Self

DEFAULT_API_BASE = "https://api.siliconflow.com"
DEFAULT_MODEL_NAME = "deepseek-ai/DeepSeek-V3.1"

_BM = TypeVar("_BM", bound=BaseModel)
_DictOrPydanticClass = Union[dict[str, Any], type[_BM], type]
_DictOrPydantic = Union[dict, _BM]


class ChatSiliconFlow(BaseChatOpenAI):
    """SiliconFlow chat model integration to access models hosted in SiliconFlow's API."""  # noqa: E501

    model_name: str = Field(alias="model", default=DEFAULT_MODEL_NAME)
    """The name of the model"""
    api_key: Optional[SecretStr] = Field(
        default_factory=secret_from_env("SILICONFLOW_API_KEY", default=None),
    )
    """SiliconFlow API key"""
    base_url: str = Field(
        default_factory=from_env("SILICONFLOW_BASE_URL", default=DEFAULT_API_BASE),
    )
    """SiliconFlow API base URL"""

    model_config = ConfigDict(populate_by_name=True)

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

    @property
    def lc_secrets(self) -> dict[str, str]:
        """A map of constructor argument names to secret ids."""
        return {"api_key": "SILICONFLOW_API_KEY"}

    def _get_ls_params(
        self,
        stop: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> LangSmithParams:
        ls_params = super()._get_ls_params(stop=stop, **kwargs)
        ls_params["ls_provider"] = "siliconflow"
        return ls_params

    @model_validator(mode="after")
    def validate_environment(self) -> Self:
        client_params: dict = {
            k: v
            for k, v in {
                "api_key": self.api_key.get_secret_value() if self.api_key else None,
                "base_url": self.base_url,
                "timeout": self.request_timeout,
                "max_retries": self.max_retries,
                "default_headers": self.default_headers,
                "default_query": self.default_query,
            }.items()
            if v is not None
        }

        if not (self.client or None):
            sync_specific: dict = {"http_client": self.http_client}
            self.root_client = openai.OpenAI(**client_params, **sync_specific)
            self.client = self.root_client.chat.completions
        if not (self.async_client or None):
            async_specific: dict = {"http_client": self.http_async_client}
            self.root_async_client = openai.AsyncOpenAI(
                **client_params,
                **async_specific,
            )
            self.async_client = self.root_async_client.chat.completions
        return self

    def _get_request_payload(
        self,
        input_: LanguageModelInput,
        *,
        stop: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> dict:
        payload = super()._get_request_payload(input_, stop=stop, **kwargs)
        for message in payload["messages"]:
            if message["role"] == "tool" and isinstance(message["content"], list):
                message["content"] = json.dumps(message["content"])
        return payload

    def _create_chat_result(
        self,
        response: Union[dict, openai.BaseModel],
        generation_info: Optional[dict] = None,
    ) -> ChatResult:
        rtn = super()._create_chat_result(response, generation_info)

        if not isinstance(response, openai.BaseModel):
            return rtn

        choices = getattr(response, "choices", None)
        if choices and hasattr(choices[0].message, "reasoning_content"):
            rtn.generations[0].message.additional_kwargs["reasoning_content"] = choices[
                0
            ].message.reasoning_content
        # Handle use via OpenRouter
        elif choices and hasattr(choices[0].message, "model_extra"):
            model_extra = choices[0].message.model_extra
            if isinstance(model_extra, dict) and (
                reasoning := model_extra.get("reasoning")
            ):
                rtn.generations[0].message.additional_kwargs["reasoning_content"] = (
                    reasoning
                )

        return rtn

    def _convert_chunk_to_generation_chunk(
        self,
        chunk: dict,
        default_chunk_class: type,
        base_generation_info: Optional[dict],
    ) -> Optional[ChatGenerationChunk]:
        generation_chunk = super()._convert_chunk_to_generation_chunk(
            chunk,
            default_chunk_class,
            base_generation_info,
        )
        if (choices := chunk.get("choices")) and generation_chunk:
            top = choices[0]
            if isinstance(generation_chunk.message, AIMessageChunk):
                if (
                    reasoning_content := top.get("delta", {}).get("reasoning_content")
                ) is not None:
                    generation_chunk.message.additional_kwargs["reasoning_content"] = (
                        reasoning_content
                    )
                # Handle use via OpenRouter
                elif (reasoning := top.get("delta", {}).get("reasoning")) is not None:
                    generation_chunk.message.additional_kwargs["reasoning_content"] = (
                        reasoning
                    )

        return generation_chunk

    def _stream(
        self,
        messages: list[BaseMessage],
        stop: Optional[list[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[ChatGenerationChunk]:
        """Stream chat completions from SiliconFlow.

        SiliconFlow API returns cumulative usage metadata in every chunk.
        LangChain aggregates these by summing them. To prevent overcounting,
        we must convert the cumulative values into deltas before yielding.
        """
        # Track the last seen cumulative totals
        last_input_tokens = 0
        last_output_tokens = 0
        try:
            for chunk in super()._stream(
                messages,
                stop=stop,
                run_manager=run_manager,
                **kwargs,
            ):
                if chunk.message.usage_metadata:
                    # 1. Get the current cumulative totals from the API response
                    current_input = chunk.message.usage_metadata.get("input_tokens", 0)
                    current_output = chunk.message.usage_metadata.get("output_tokens", 0)

                    # 2. Calculate the delta (New - Old)
                    # This effectively "strips" the already-counted tokens from this chunk
                    input_delta = current_input - last_input_tokens
                    output_delta = current_output - last_output_tokens

                    # 3. Update the chunk with the calculated deltas
                    chunk.message.usage_metadata["input_tokens"] = input_delta
                    chunk.message.usage_metadata["output_tokens"] = output_delta
                    chunk.message.usage_metadata["total_tokens"] = input_delta + output_delta

                    # 4. Update state for the next iteration
                    last_input_tokens = current_input
                    last_output_tokens = current_output

                yield chunk
        except JSONDecodeError as e:
            msg = (
                "SiliconFlow API returned an invalid response. "
                "Please check the API status and try again."
            )
            raise JSONDecodeError(
                msg,
                e.doc,
                e.pos,
            ) from e

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: Optional[list[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        try:
            return super()._generate(
                messages,
                stop=stop,
                run_manager=run_manager,
                **kwargs,
            )
        except JSONDecodeError as e:
            msg = (
                "SiliconFlow API returned an invalid response. "
                "Please check the API status and try again."
            )
            raise JSONDecodeError(
                msg,
                e.doc,
                e.pos,
            ) from e

    async def _astream(
        self,
        messages: list[BaseMessage],
        stop: Optional[list[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> AsyncIterator[ChatGenerationChunk]:
        """Async stream chat completions from SiliconFlow.

        Applies the same delta logic as _stream to prevent overcounting usage.
        """
        last_input_tokens = 0
        last_output_tokens = 0
        try:
            async for chunk in super()._astream(
                messages,
                stop=stop,
                run_manager=run_manager,
                **kwargs,
            ):
                if chunk.message.usage_metadata:
                    current_input = chunk.message.usage_metadata.get("input_tokens", 0)
                    current_output = chunk.message.usage_metadata.get("output_tokens", 0)

                    input_delta = current_input - last_input_tokens
                    output_delta = current_output - last_output_tokens

                    chunk.message.usage_metadata["input_tokens"] = input_delta
                    chunk.message.usage_metadata["output_tokens"] = output_delta
                    chunk.message.usage_metadata["total_tokens"] = input_delta + output_delta

                    last_input_tokens = current_input
                    last_output_tokens = current_output

                yield chunk
        except JSONDecodeError as e:
            msg = "SiliconFlow API returned an invalid response."
            raise JSONDecodeError(msg, e.doc, e.pos) from e
