"""
OpenCode Style Event Adapter
============================

将 LangGraph `astream_events` 的原始事件流，转换为 OpenCode 风格的统一事件。

核心目标：
- 保持 LangGraph 流式特性（AsyncIterator）
- 按文档中定义的事件语义进行 1:1 / 1:N 映射
- 提供简单、前端友好的结构化事件
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, AsyncIterator, Optional, Dict, AsyncGenerator


class OpenCodeEventType(Enum):
    # 会话级别
    START = "start"
    FINISH = "finish"

    # 推理/思考
    REASONING_START = "reasoning-start"
    REASONING_DELTA = "reasoning-delta"
    REASONING_END = "reasoning-end"

    # 文字生成
    TEXT_START = "text-start"
    TEXT_DELTA = "text-delta"
    TEXT_END = "text-end"

    # 工具调用
    TOOL_INPUT = "tool-input"
    TOOL_INPUT_START = "tool-input-start"
    TOOL_INPUT_DELTA = "tool-input-delta"
    TOOL_CALL = "tool-call"
    TOOL_RESULT = "tool-result"
    TOOL_ERROR = "tool-error"

    # 步骤/循环
    START_STEP = "start-step"
    FINISH_STEP = "finish-step"

    # 错误
    ERROR = "error"


@dataclass
class OpenCodeEvent:
    """OpenCode 风格事件"""

    event_type: OpenCodeEventType
    data: Dict[str, Any]
    step_index: Optional[int] = None
    tool_name: Optional[str] = None
    timestamp: float = field(default_factory=lambda: time.time())


@dataclass
class AdapterState:
    """适配器状态"""

    # 图级别
    graph_name: str = ""
    is_graph_started: bool = False

    # 步骤级别
    step_index: int = 0
    is_step_started: bool = False
    step_start_node: str = ""

    # 模型级别
    is_model_started: bool = False
    is_reasoning_active: bool = False
    is_text_active: bool = False
    has_reasoning: bool = False  # 是否是推理模型

    # 工具级别
    current_tool: Optional[str] = None
    tool_input_chunks: list[str] = field(default_factory=list)

    # 本轮答案累积（用于在 finish 事件中携带完整 answer；优先 text，其次 reason，最后 on_chain_end）
    accumulated_answer: str = ""   # 仅累积 text（最后一轮 model）
    accumulated_reasoning: str = ""  # 仅累积 reasoning（最后一轮 model，CoT 无 text 时作 answer）

    # 辅助判断
    last_chain_end_name: str = ""


class OpenCodeStyleEventAdapter:
    """OpenCode 风格事件适配器（基于 LangGraph astream_events）"""

    def __init__(self, graph_name: str = "LangGraph", config: Optional[dict] = None) -> None:
        self.graph_name = graph_name
        self.state = AdapterState()

        config = config or {}

        # 步骤边界配置
        self.loop_entry_nodes = set(config.get("loop_entry_nodes", {"model"}))
        self.loop_end_nodes = set(
            config.get("loop_end_nodes", {"prune_and_compress", "loop_controller"})
        )

        # 开关配置
        self.enable_reasoning_events: bool = config.get("enable_reasoning_events", True)
        self.enable_step_events: bool = config.get("enable_step_events", True)
        self.synthesize_tool_call: bool = config.get("synthesize_tool_call", True)
        self.enable_tool_input_start: bool = config.get("enable_tool_input_start", False)

    async def adapt(
        self, events: AsyncIterator[Dict[str, Any]]
    ) -> AsyncIterator[OpenCodeEvent]:
        """主转换函数：将 LangGraph 事件流转换为 OpenCode 风格事件流。"""

        async for event in events:
            event_type = event.get("event")
            name = event.get("name", "") or ""
            data = event.get("data", {}) or {}

            if event_type == "on_chain_start":
                async for oc_event in self._handle_chain_start(name, data):
                    yield oc_event

            elif event_type == "on_chain_end":
                async for oc_event in self._handle_chain_end(name, data):
                    yield oc_event

            elif event_type == "on_chat_model_start":
                async for oc_event in self._handle_model_start(name, data):
                    if oc_event is not None:
                        yield oc_event

            elif event_type == "on_chat_model_stream":
                async for oc_event in self._handle_model_stream(name, data):
                    if oc_event is not None:
                        yield oc_event

            elif event_type == "on_chat_model_end":
                async for oc_event in self._handle_model_end(name, data):
                    if oc_event is not None:
                        yield oc_event

            elif event_type == "on_tool_start":
                async for oc_event in self._handle_tool_start(name, data):
                    yield oc_event

            elif event_type == "on_tool_end":
                async for oc_event in self._handle_tool_end(name, data):
                    yield oc_event

            elif event_type == "on_chain_error":
                async for oc_event in self._handle_error(name, data):
                    yield oc_event

    def _reset_state(self) -> None:
        """重置内部状态（会话结束后调用）。"""
        self.state = AdapterState()

    # ===== Chain 事件处理 =====

    async def _handle_chain_start(
        self, name: str, data: Dict[str, Any]
    ) -> AsyncGenerator[OpenCodeEvent, None]:
        """处理 on_chain_start"""

        # 情况 1：图开始
        if not self.state.is_graph_started:
            self.state.graph_name = name or self.graph_name
            self.state.is_graph_started = True
            yield OpenCodeEvent(
                event_type=OpenCodeEventType.START,
                data={"graph_name": self.state.graph_name},
            )
            return

        # 情况 2：步骤开始（loop_entry_node）
        #
        # 严格成对策略（与 factory.py 的循环语义对齐）：
        # - 一个循环（step）定义为：从一次进入 loop_entry_node 开始，
        #   到下一次进入 loop_entry_node 之前（或图结束）为止。
        # - 因此，finish-step 不应依赖 loop_controller / prune_and_compress 等中间节点，
        #   而应在“下一次 loop_entry_node 发生前”或“图结束时”补齐。
        if self.enable_step_events and self._is_loop_entry_node(name):
            # 如果上一个 step 仍处于打开状态，先在本次 step 开始前关闭它
            if self.state.is_step_started and self.state.step_index > 0:
                yield OpenCodeEvent(
                    event_type=OpenCodeEventType.FINISH_STEP,
                    data={"step_index": self.state.step_index},
                    step_index=self.state.step_index,
                )
                self.state.is_step_started = False

            self.state.step_index += 1
            self.state.is_step_started = True
            self.state.step_start_node = name
            # 每进入新 step 清空累积，只保留「最后一轮」model 的 text + reason
            self.state.accumulated_answer = ""
            self.state.accumulated_reasoning = ""

            yield OpenCodeEvent(
                event_type=OpenCodeEventType.START_STEP,
                data={"step_index": self.state.step_index},
                step_index=self.state.step_index,
            )

    def _is_loop_entry_node(self, name: str) -> bool:
        """判断是否是循环入口节点"""
        if name in self.loop_entry_nodes:
            return True
        if name.endswith(".before_model"):
            return True
        return False

    def _extract_answer_from_chain_end_output(self, data: Dict[str, Any]) -> str:
        """从 on_chain_end 的 output 中取最后一条 AIMessage.content。"""
        output = data.get("output") if isinstance(data, dict) else None
        if not isinstance(output, dict) or "messages" not in output:
            return ""
        messages = output.get("messages")
        if not isinstance(messages, (list, tuple)):
            return ""
        for msg in reversed(messages):
            if getattr(msg, "__class__", None) and getattr(msg.__class__, "__name__", "") == "AIMessage":
                content = getattr(msg, "content", None)
                if isinstance(content, str):
                    return content
        return ""

    async def _handle_chain_end(
        self, name: str, data: Dict[str, Any]
    ) -> AsyncGenerator[OpenCodeEvent, None]:
        """处理 on_chain_end"""

        # 图结束：如果仍有打开的 step，先补齐 finish-step，再发 finish
        if name == self.state.graph_name:
            if self.enable_step_events and self.state.is_step_started and self.state.step_index > 0:
                yield OpenCodeEvent(
                    event_type=OpenCodeEventType.FINISH_STEP,
                    data={"step_index": self.state.step_index},
                    step_index=self.state.step_index,
                )
                self.state.is_step_started = False

            # answer：优先 text，其次 reason（CoT 无 text 时），最后 on_chain_end 兜底
            answer = self.state.accumulated_answer.strip()
            if not answer:
                answer = self.state.accumulated_reasoning.strip()
            if not answer:
                answer = self._extract_answer_from_chain_end_output(data)
            yield OpenCodeEvent(
                event_type=OpenCodeEventType.FINISH,
                data={
                    "graph_name": name,
                    "answer": answer,
                },
            )
            self._reset_state()

    def _is_step_end_node(self, name: str) -> bool:
        """判断是否是步骤结束节点"""
        return name in self.loop_end_nodes

    # ===== Model 事件处理 =====

    async def _handle_model_start(
        self, name: str, data: Dict[str, Any]
    ) -> AsyncGenerator[Optional[OpenCodeEvent], None]:
        """处理 on_chat_model_start"""
        self.state.is_model_started = True
        # 每次 model 开始时清空累积，只保留最后一次 model 的 text + reason
        self.state.accumulated_answer = ""
        self.state.accumulated_reasoning = ""
        # 暂不生成事件，等待 stream 时判断是否有 reasoning/text
        if False:  # pragma: no cover - 仅用于保持 async generator 形式
            yield None

    async def _handle_model_stream(
        self, name: str, data: Dict[str, Any]
    ) -> AsyncGenerator[Optional[OpenCodeEvent], None]:
        """处理 on_chat_model_stream"""
        chunk = data.get("chunk")
        if not chunk:
            return

        # ===== 处理 reasoning / text（对齐 LangChain content_blocks + additional_kwargs 语义）=====

        reasoning_segments: list[str] = []
        text_segments: list[str] = []

        # 1. 优先解析 content_blocks（LangChain v1 标准）
        try:
            blocks = None
            if hasattr(chunk, "content_blocks") and getattr(chunk, "content_blocks"):
                blocks = getattr(chunk, "content_blocks")
            elif isinstance(chunk, dict) and chunk.get("content_blocks"):
                blocks = chunk["content_blocks"]

            if blocks and isinstance(blocks, (list, tuple)):
                for block in blocks:
                    if not isinstance(block, dict):
                        continue
                    btype = block.get("type", "")
                    if btype in ("reasoning", "thinking"):
                        r = block.get("reasoning") or block.get("thinking") or ""
                        if isinstance(r, str) and r:
                            reasoning_segments.append(r)
                    elif btype == "text":
                        t = block.get("text", "")
                        if isinstance(t, str) and t:
                            text_segments.append(t)
        except Exception:
            # content_blocks 解析失败则退回后续路径
            pass

        # 2. 若没有从 content_blocks 中拿到 reasoning，则补充检查 additional_kwargs
        #    - reasoning_content (str): Ollama, DeepSeek, XAI, Groq 等
        #    - reasoning (dict): OpenAI Responses 风格结构
        if not reasoning_segments and hasattr(chunk, "additional_kwargs"):
            ak = getattr(chunk, "additional_kwargs") or {}
            if isinstance(ak, dict):
                rc = ak.get("reasoning_content")
                if isinstance(rc, str) and rc:
                    reasoning_segments.append(rc)
                else:
                    robj = ak.get("reasoning")
                    if isinstance(robj, dict):
                        r = robj.get("reasoning")
                        if isinstance(r, str) and r:
                            reasoning_segments.append(r)
                        else:
                            summary = robj.get("summary")
                            if isinstance(summary, list):
                                for part in summary:
                                    if isinstance(part, dict) and isinstance(
                                        part.get("text"), str
                                    ):
                                        reasoning_segments.append(part["text"])

        # 3. 兼容 chunk.reasoning / chunk.content（如部分自定义 ChatModel）
        if self.enable_reasoning_events and not reasoning_segments:
            if hasattr(chunk, "reasoning") and getattr(chunk, "reasoning"):
                r = getattr(chunk, "reasoning")
                if isinstance(r, str) and r:
                    reasoning_segments.append(r)

        if not text_segments and hasattr(chunk, "content"):
            c = getattr(chunk, "content")
            if isinstance(c, str) and c:
                text_segments.append(c)

        # ===== 发出 OpenCode 风格 reasoning-* / text-* 事件 =====

        # reasoning 流（推理模型）
        if self.enable_reasoning_events and reasoning_segments:
            combined_reasoning = "".join(reasoning_segments)
            if combined_reasoning:
                if not self.state.is_reasoning_active:
                    self.state.is_reasoning_active = True
                    self.state.has_reasoning = True
                    yield OpenCodeEvent(
                        event_type=OpenCodeEventType.REASONING_START,
                        data={},
                        step_index=self.state.step_index,
                    )

                yield OpenCodeEvent(
                    event_type=OpenCodeEventType.REASONING_DELTA,
                    data={"text": combined_reasoning},
                    step_index=self.state.step_index,
                )
                self.state.accumulated_reasoning += combined_reasoning

        # 文本流（所有模型）
        if text_segments:
            combined_text = "".join(text_segments)
            if combined_text:
                if not self.state.is_text_active:
                    # 推理模型：reasoning 结束后才开始 text
                    if self.enable_reasoning_events and self.state.is_reasoning_active:
                        yield OpenCodeEvent(
                            event_type=OpenCodeEventType.REASONING_END,
                            data={},
                            step_index=self.state.step_index,
                        )
                        self.state.is_reasoning_active = False

                    self.state.is_text_active = True
                    yield OpenCodeEvent(
                        event_type=OpenCodeEventType.TEXT_START,
                        data={},
                        step_index=self.state.step_index,
                    )

                yield OpenCodeEvent(
                    event_type=OpenCodeEventType.TEXT_DELTA,
                    data={"text": combined_text},
                    step_index=self.state.step_index,
                )
                self.state.accumulated_answer += combined_text

    async def _handle_model_end(
        self, name: str, data: Dict[str, Any]
    ) -> AsyncGenerator[Optional[OpenCodeEvent], None]:
        """处理 on_chat_model_end"""

        # 结束 reasoning（如果还在活跃）
        if self.enable_reasoning_events and self.state.is_reasoning_active:
            yield OpenCodeEvent(
                event_type=OpenCodeEventType.REASONING_END,
                data={},
                step_index=self.state.step_index,
            )
            self.state.is_reasoning_active = False

        # 结束 text（如果还在活跃）
        if self.state.is_text_active:
            yield OpenCodeEvent(
                event_type=OpenCodeEventType.TEXT_END,
                data={},
                step_index=self.state.step_index,
            )
            self.state.is_text_active = False

        self.state.is_model_started = False

    # ===== Tool 事件处理 =====

    async def _handle_tool_start(
        self, name: str, data: Dict[str, Any]
    ) -> AsyncGenerator[OpenCodeEvent, None]:
        """处理 on_tool_start - 一对多映射"""

        self.state.current_tool = name

        # 可选：tool-input-start
        if self.enable_tool_input_start:
            yield OpenCodeEvent(
                event_type=OpenCodeEventType.TOOL_INPUT_START,
                data={"tool_name": name},
                step_index=self.state.step_index,
                tool_name=name,
            )

        # 1. tool-input：完整的工具输入（归一化 input 为 dict，便于前端展示参数）
        raw_input = data.get("input", data.get("inputs", {}))
        if isinstance(raw_input, str):
            try:
                tool_input = json.loads(raw_input) if (raw_input and raw_input.strip()) else {}
            except Exception:
                tool_input = {}
        elif isinstance(raw_input, dict):
            tool_input = raw_input
        else:
            tool_input = {}
        yield OpenCodeEvent(
            event_type=OpenCodeEventType.TOOL_INPUT,
            data={
                "tool_name": name,
                "input": tool_input,
                "params": tool_input,  # 前端兼容：同时提供 params 便于 getToolDisplayText 读取
            },
            step_index=self.state.step_index,
            tool_name=name,
        )

        # 2. tool-input-delta：LangGraph 不支持，跳过（仅保留类型以兼容枚举）

        # 3. tool-call：合成事件（表示工具正在执行）
        if self.synthesize_tool_call:
            yield OpenCodeEvent(
                event_type=OpenCodeEventType.TOOL_CALL,
                data={
                    "tool_name": name,
                    "status": "running",
                },
                step_index=self.state.step_index,
                tool_name=name,
            )

    async def _handle_tool_end(
        self, name: str, data: Dict[str, Any]
    ) -> AsyncGenerator[OpenCodeEvent, None]:
        """处理 on_tool_end - 条件映射"""

        output = data.get("output")

        # 判断是否是错误
        is_error = self._is_tool_error(output)

        if is_error:
            yield OpenCodeEvent(
                event_type=OpenCodeEventType.TOOL_ERROR,
                data={
                    "tool_name": name,
                    "error": str(output),
                    "error_type": self._get_error_type(output),
                    "status": "blocked" if isinstance(output, PermissionError) else "error",
                },
                step_index=self.state.step_index,
                tool_name=name,
            )
        else:
            # 保证 output 可 JSON 序列化：ToolMessage 等只下发 .content
            out_val = output
            if hasattr(output, "content"):
                out_val = getattr(output, "content")
            yield OpenCodeEvent(
                event_type=OpenCodeEventType.TOOL_RESULT,
                data={
                    "tool_name": name,
                    "output": out_val,
                },
                step_index=self.state.step_index,
                tool_name=name,
            )

        self.state.current_tool = None

    def _is_tool_error(self, output: Any) -> bool:
        """判断工具输出是否是错误"""
        if output is None:
            return False
        if isinstance(output, Exception):
            return True
        if isinstance(output, dict):
            return bool(output.get("error") is not None or output.get("status") == "error")
        return False

    def _get_error_type(self, output: Any) -> str:
        """获取错误类型"""
        if isinstance(output, PermissionError):
            return "permission_denied"
        if isinstance(output, TimeoutError):
            return "timeout"
        if isinstance(output, ValueError):
            return "invalid_input"
        return "unknown"

    # ===== Error 事件处理 =====

    async def _handle_error(
        self, name: str, data: Dict[str, Any]
    ) -> AsyncGenerator[OpenCodeEvent, None]:
        """处理 on_chain_error"""

        error = data.get("error", "Unknown error")

        yield OpenCodeEvent(
            event_type=OpenCodeEventType.ERROR,
            data={
                "error": str(error),
                "error_type": type(error).__name__ if isinstance(error, Exception) else "Unknown",
                "node_name": name,
            },
            step_index=self.state.step_index,
        )


__all__ = [
    "OpenCodeEventType",
    "OpenCodeEvent",
    "AdapterState",
    "OpenCodeStyleEventAdapter",
]

