"""Agent Loop 引擎 — Thinking+JSON 推理循环，pluggy 驱动 hook。

算法保真 #11: 完整保留 TRM4 core/loop.py 逻辑:
- json_repair 兜底解析
- submit_answer 终止
- 无效工具（ValueError）不计步
- pluggy hook 生命周期（before_step / after_tool / after_step / on_finish）

TRM4 → TRM5 有意变更（非简化）:
- 同步 → 全异步（async/await）
- client: Any → llm: LLMProvider（Protocol 类型化）
- tool_fn: Callable → ToolDispatcher.dispatch()（Protocol + context）
- Step 新增 call_id（从 LLMResponse.call_id 透传）
- thinking 从 getattr(msg, "reasoning_content") → response.thinking（adapters 已统一剥离）
- token 用量从 response.usage.prompt_tokens → response.prompt_tokens（LLMResponse 扁平化）
"""

from __future__ import annotations

import asyncio
import json
import re
from typing import TYPE_CHECKING, Any

import pluggy
from json_repair import repair_json
from loguru import logger

from core.agent.protocols import AgentLoopSpec, ToolDispatcher
from core.agent.types import LoopResult, Step

if TYPE_CHECKING:
    from core.protocols import LLMProvider
    from core.types import LLMResponse

# deepseek 等模型稳定输出变体：```json 围栏包裹 JSON 体
_CODE_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*\n?|\n?\s*```\s*$")


async def _call_hook(hook: Any, **kwargs: Any) -> list[Any]:
    """调用 pluggy hook 并 await 异步返回值。

    pluggy 本身是同步调度，但 hookimpl 可以是 async def，
    此时 hook() 返回 coroutine 列表，需要逐个 await。

    参数:
        hook: pluggy hook caller（如 pm.hook.before_step）。
        **kwargs: 传递给 hook 的关键字参数。

    返回:
        已 resolve 的返回值列表。
    """
    results = hook(**kwargs)
    if results is not None:
        resolved = []
        for r in results:
            if hasattr(r, "__await__"):
                resolved.append(await r)
            else:
                resolved.append(r)
        return resolved
    return []


class AgentLoop:
    """Thinking+JSON 推理循环引擎。

    类比 nn.Module: 接收 prompt + 工具调度器，返回 LoopResult。
    不感知视频树、QA、数据库等领域概念。

    参数:
        llm: LLMProvider 实例（Protocol 类型化，提供 async chat 方法）。
        max_steps: 最大有效步数（每次成功工具调用计一步）。
        max_retries: JSON 解析连续失败的最大容忍次数。
        step_retries: LLM 瞬时异常的步级重试次数（不含首次调用）。
        step_retry_delays: 步级重试的退避秒数序列，超出部分取末值；
            不得为空（空序列构造时抛 ValueError，fail-fast）。
        retryable_exceptions: 可重试异常元组，默认 (TimeoutError, OSError)——
            ssl.SSLError/ConnectionError 均为 OSError 子类，覆盖穿透
            GovernedLLMClient 内部重试栈的瞬时异常；openai API 类异常由
            治理层负责，core 不依赖 SDK。
    """

    def __init__(
        self,
        llm: LLMProvider,
        max_steps: int,
        max_retries: int = 3,
        *,
        step_retries: int = 2,
        step_retry_delays: tuple[float, ...] = (20.0, 40.0),
        retryable_exceptions: tuple[type[BaseException], ...] = (TimeoutError, OSError),
    ) -> None:
        if not step_retry_delays:
            raise ValueError("step_retry_delays 不得为空")
        self._llm = llm
        self._max_steps = max_steps
        self._max_retries = max_retries
        self._step_retries = step_retries
        self._step_retry_delays = step_retry_delays
        self._retryable_exceptions = retryable_exceptions

    async def run(
        self,
        system_prompt: str,
        user_prompt: str,
        tool_dispatcher: ToolDispatcher,
        plugins: list[object] | None = None,
        *,
        session_id: str | None = None,
        cache_salt: str | None = None,
    ) -> LoopResult:
        """执行 Thinking+JSON 推理循环。

        参数:
            system_prompt: 系统提示词。
            user_prompt: 用户提示词。
            tool_dispatcher: 工具调度器，ToolDispatcher Protocol 实例。
            plugins: pluggy 插件列表。
            session_id: 会话 ID，透传给 LLMProvider。
            cache_salt: 缓存盐，透传给 LLMProvider（如训练用 run_id 跨 epoch 重采样）。

        返回:
            LoopResult 实例，包含推理步骤、token 用量、终止原因。
        """
        pm = self._create_plugin_manager(plugins)
        messages: list[dict[str, Any]] = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
        steps: list[Step] = []
        token_usage: dict[str, int] = {"prompt_tokens": 0, "completion_tokens": 0}
        step_count = 0
        retry_count = 0
        iteration = 0

        while step_count < self._max_steps:
            await _call_hook(pm.hook.before_step, iteration=iteration, messages=messages)

            # Phase 1: LLM 调用（步级重试：防穿透 GovernedLLMClient 的瞬时异常）
            try:
                response = await self._call_llm_with_step_retry(
                    messages, token_usage, session_id=session_id, cache_salt=cache_salt
                )
            except Exception as e:
                logger.error("LLM API 调用失败（{}）: {}", type(e).__name__, e)
                result = LoopResult(
                    steps=steps,
                    steps_used=step_count,
                    token_usage=token_usage,
                    stop_reason="error",
                )
                await _call_hook(pm.hook.on_finish, result=result)
                return result

            # Phase 2: 解析响应
            parsed = self._parse_response(response)
            if parsed is None:
                retry_count += 1
                logger.warning("响应解析失败 (retry {}/{})", retry_count, self._max_retries)
                messages.append({"role": "assistant", "content": response.content})
                messages.append(
                    {
                        "role": "user",
                        "content": (
                            "你的输出不是合法 JSON。请严格输出 JSON 格式："
                            '{"reflect": {...}, "plan": {...}, '
                            '"action": {"tool": "...", "args": {...}}}'
                        ),
                    }
                )
                if retry_count >= self._max_retries:
                    result = LoopResult(
                        steps=steps,
                        steps_used=step_count,
                        token_usage=token_usage,
                        stop_reason="parse_error",
                    )
                    await _call_hook(pm.hook.after_step, iteration=iteration, messages=messages)
                    await _call_hook(pm.hook.on_finish, result=result)
                    return result
                await _call_hook(pm.hook.after_step, iteration=iteration, messages=messages)
                iteration += 1
                continue

            thought, reflect, plan, raw_content, action, call_id = parsed
            retry_count = 0
            messages.append({"role": "assistant", "content": raw_content})

            # Phase 3: 执行工具
            tool_name: str = action["tool"]
            tool_args: dict[str, Any] = action["args"]
            context: dict[str, Any] = {
                "session_id": session_id,
                "iteration": iteration,
            }
            output, is_valid = await self._execute_tool(
                tool_dispatcher, tool_name, tool_args, context=context
            )
            if not is_valid:
                messages.append(
                    {
                        "role": "user",
                        "content": f"[工具调用无效: {tool_name}] {output}",
                    }
                )
                await _call_hook(pm.hook.after_step, iteration=iteration, messages=messages)
                iteration += 1
                continue

            step_count += 1
            step = Step(
                thought=thought,
                reflect=reflect,
                plan=plan,
                tool_call={"tool": tool_name, "args": tool_args},
                tool_output=output,
                raw_content=raw_content,
                call_id=call_id,
            )
            steps.append(step)

            # Phase 4: Hook + 反馈
            hints = await _call_hook(pm.hook.after_tool, iteration=iteration, step=step)
            feedback = self._build_feedback(tool_name, output, hints)
            messages.append(feedback)
            await _call_hook(pm.hook.after_step, iteration=iteration, messages=messages)

            # Phase 5: 终止检查
            if tool_name == "submit_answer":
                result = LoopResult(
                    result=tool_args,
                    steps=steps,
                    steps_used=step_count,
                    token_usage=token_usage,
                    stop_reason="finished",
                )
                await _call_hook(pm.hook.on_finish, result=result)
                return result

            iteration += 1

        # 预算耗尽
        result = LoopResult(
            steps=steps,
            steps_used=step_count,
            token_usage=token_usage,
            stop_reason="budget_exceeded",
        )
        await _call_hook(pm.hook.on_finish, result=result)
        return result

    def _create_plugin_manager(self, plugins: list[object] | None) -> pluggy.PluginManager:
        """创建并注册 plugins 的 PluginManager。

        参数:
            plugins: pluggy 插件列表，可为 None。

        返回:
            配置好的 PluginManager 实例。
        """
        pm = pluggy.PluginManager("agent_loop")
        pm.add_hookspecs(AgentLoopSpec)
        for plugin in plugins or []:
            pm.register(plugin)
        return pm

    async def _call_llm_with_step_retry(
        self,
        messages: list[dict[str, Any]],
        token_usage: dict[str, int],
        *,
        session_id: str | None = None,
        cache_salt: str | None = None,
    ) -> LLMResponse:
        """带步级重试的 LLM 调用，兜底穿透治理层重试栈的瞬时异常。

        可重试异常（self._retryable_exceptions）按 self._step_retry_delays
        退避后重发，重试预算（self._step_retries）耗尽后抛出最后一次异常；
        非可重试异常直接穿透（fail-fast）。asyncio.CancelledError 继承
        BaseException，天然不被捕获——取消信号照常传播。失败尝试的 error
        遥测由 GovernedLLMClient 内部负责，此处仅 loguru 记录。

        参数:
            messages: 消息历史。
            token_usage: 可变字典，就地累加。
            session_id: 会话 ID，透传给 LLMProvider。

        返回:
            LLMResponse 实例。

        异常:
            重试耗尽或非可重试时，原样抛出底层异常。
        """
        step_attempt = 0
        while True:
            try:
                return await self._call_llm(
                    messages, token_usage, session_id=session_id, cache_salt=cache_salt
                )
            except self._retryable_exceptions as e:
                step_attempt += 1
                if step_attempt > self._step_retries:
                    raise
                delay = self._step_retry_delays[
                    min(step_attempt - 1, len(self._step_retry_delays) - 1)
                ]
                logger.warning(
                    "LLM 瞬时异常，步级重试 {}/{}（{}s 后重发）: {}",
                    step_attempt,
                    self._step_retries,
                    delay,
                    e,
                )
                await asyncio.sleep(delay)

    async def _call_llm(
        self,
        messages: list[dict[str, Any]],
        token_usage: dict[str, int],
        *,
        session_id: str | None = None,
        cache_salt: str | None = None,
    ) -> LLMResponse:
        """调用 LLMProvider 并累加 token 使用量。

        参数:
            messages: 消息历史。
            token_usage: 可变字典，就地累加。
            session_id: 会话 ID，透传给 LLMProvider。
            cache_salt: 缓存盐，透传给 LLMProvider（跨 epoch 重采样）。

        返回:
            LLMResponse 实例。
        """
        response = await self._llm.chat(messages, session_id=session_id, cache_salt=cache_salt)
        token_usage["prompt_tokens"] += response.prompt_tokens
        token_usage["completion_tokens"] += response.completion_tokens
        return response

    def _parse_response(
        self, response: LLMResponse
    ) -> tuple[str, dict, dict, str, dict, str] | None:
        """从 LLMResponse 中提取结构化决策数据。

        解析流程: content → 剥除 ```json 围栏 → repair_json → json.loads
        → 收拢 action 平铺参数 → 校验 action/tool/args。

        参数:
            response: LLMResponse 实例。

        返回:
            解析成功返回 (thought, reflect, plan, raw_content, action, call_id)；
            解析失败返回 None。
        """
        content = response.content
        thought = response.thinking

        if not content.strip():
            return None

        repaired = repair_json(_CODE_FENCE_RE.sub("", content).strip())
        try:
            data = json.loads(repaired)
        except (json.JSONDecodeError, ValueError):
            return None

        if not isinstance(data, dict) or "action" not in data:
            return None

        action = self._normalize_action(data["action"])
        if not isinstance(action, dict) or "tool" not in action or "args" not in action:
            return None

        reflect = data.get("reflect", {})
        plan = data.get("plan", {})
        return thought, reflect, plan, content, action, response.call_id

    @staticmethod
    def _normalize_action(action: Any) -> Any:
        """收拢 deepseek 变体的 action 平铺参数。

        deepseek 等模型稳定输出变体: 工具参数平铺在 action 下（缺 args
        嵌套），确定性收拢为标准 {"tool": ..., "args": {...}} 结构。
        仅当除 tool 外至少存在一个平铺参数键时才收拢；无参结构
        （如 {"tool": "x"}）原样返回交由调用方校验拒绝，避免把缺参
        错误静默升级为空 args 合法结构。标准嵌套与非法结构同样原样返回。

        参数:
            action: 从 LLM 输出解析出的 action 字段（任意类型）。

        返回:
            归一化后的 action（仅带平铺参数的变体被改写，其余原样返回）。
        """
        if isinstance(action, dict) and "tool" in action and "args" not in action:
            flat_args = {k: v for k, v in action.items() if k != "tool"}
            if flat_args:
                return {"tool": action["tool"], "args": flat_args}
        return action

    async def _execute_tool(
        self,
        dispatcher: ToolDispatcher,
        name: str,
        args: dict[str, Any],
        *,
        context: dict[str, Any],
    ) -> tuple[str, bool]:
        """执行工具调用。

        参数:
            dispatcher: 工具调度器。
            name: 工具名称。
            args: 工具参数。
            context: 调用上下文（session_id、iteration 等）。

        返回:
            (output, is_valid) — ValueError 时 is_valid=False 且不计步数。
        """
        try:
            output = await dispatcher.dispatch(name, args, context=context)
            return output, True
        except ValueError as e:
            return f"工具调用失败: {e}", False

    def _build_feedback(
        self,
        tool_name: str,
        tool_output: str,
        hints: list[str | None],
    ) -> dict[str, Any]:
        """组装工具结果反馈消息。

        参数:
            tool_name: 工具名称。
            tool_output: 工具原始输出。
            hints: hook 返回的 hint 列表（含 None）。

        返回:
            user role 消息字典。
        """
        parts = [f"[工具执行结果: {tool_name}]", tool_output]
        for hint in hints:
            if hint is not None:
                parts.append(hint)
        return {"role": "user", "content": "\n".join(parts)}