"""
CompatibleChatOpenAI
====================

为 **任意 OpenAI 兼容聊天服务** 提供一个轻量级的 LangChain ChatModel 封装，
在不 fork `langchain-openai` 的前提下：

- 复用官方 `ChatOpenAI` 的绝大部分行为（重试、流式、usage 统计等）
- 额外把流式响应中的 `reasoning_content` 注入到
  `AIMessageChunk.additional_kwargs["reasoning_content"]`，让下游组件按
  LangChain v2 标准路径解析 reasoning。
"""

from __future__ import annotations

from typing import Any, Dict, Type

from langchain_openai import ChatOpenAI
from langchain_core.messages import AIMessageChunk, BaseMessageChunk
from langchain_core.outputs import ChatGenerationChunk


class CompatibleChatOpenAI(ChatOpenAI):
    """
    OpenAI 兼容 ChatModel 封装（支持 reasoning_content）。

    使用方式（示例）::

        from langchain_agentx.provider import CompatibleChatOpenAI

        llm = CompatibleChatOpenAI(
            model="your-openai-compatible-model",
            api_key="你的 key",
            base_url="https://your-openai-compatible-endpoint/v1",
            streaming=True,
        )

        # 之后可以像 ChatOpenAI 一样使用:
        # result = await llm.ainvoke([...])
        # async for chunk in llm.astream([...]):
        #     ...
    """

    def _convert_chunk_to_generation_chunk(
        self,
        chunk: Dict[str, Any],
        default_chunk_class: Type[BaseMessageChunk],
        base_generation_info: Dict[str, Any] | None,
    ) -> ChatGenerationChunk | None:
        """
        在父类基础上，额外抽取 OpenAI 兼容服务的 `reasoning_content` 字段。

        上游 OpenAI 兼容流式 chunk 结构（简化）::

            {
                "id": "...",
                "choices": [{
                    "delta": {
                        "role": "assistant",
                        "content": "...",
                        "reasoning_content": "..."   # COT 思考过程
                    },
                    "finish_reason": null,
                    "index": 0,
                }],
                "model": "xxx",
                ...
            }
        """
        # 先让 ChatOpenAI 做标准转换
        gen_chunk = super()._convert_chunk_to_generation_chunk(
            chunk, default_chunk_class, base_generation_info
        )
        if gen_chunk is None:
            return None

        message = gen_chunk.message
        # 仅对 AIMessageChunk 追加 reasoning_content
        if not isinstance(message, AIMessageChunk):
            return gen_chunk

        # 从原始 chunk 中取出 delta.reasoning_content（OpenAI 兼容服务的自定义字段）
        choices = chunk.get("choices", []) or chunk.get("chunk", {}).get("choices", [])
        if not choices:
            return gen_chunk

        delta = choices[0].get("delta") or {}
        rc = delta.get("reasoning_content")
        if isinstance(rc, str) and rc:
            # 将 reasoning_content 写入 additional_kwargs
            ak = getattr(message, "additional_kwargs", None) or {}
            ak["reasoning_content"] = rc
            message.additional_kwargs = ak

        return gen_chunk


__all__ = ["CompatibleChatOpenAI"]

