"""子智能体生命周期跟踪。

## 为什么换掉 1.5 的做法

1.5 的 ``_SubAgentTracker`` 从子图命名空间字符串 ``"tools:<X>"`` 里**反查
``subagent_type``**：拿到 ``<X>`` 这个 task tool_call 的 id 后，还得回头翻
历史 tool_calls 才能知道它是哪个子智能体。这条链路脆在两个地方——命名空间
格式由框架内部决定（升级可能变），以及反查依赖「id 还在最近的历史窗口里」
（一旦被上下文压缩裁掉，反查就落空，子智能体身份丢失）。

## 2.0 的做法

在**收到 task 工具调用的那一刻**就把 ``tool_call_id → {subagent_type, 描述}``
写进一张表，之后所有子图事件只用 id 查表：

- 不解析命名空间语义，只取 ``ns[0]`` 里 ``tools:`` 后面的 id
- 不依赖历史窗口，表由本对象持有，生命周期覆盖整轮请求
- 表项在 task 结果回来时标记 done，可据此判断「还有几个子智能体在跑」

命名空间格式本身仍来自框架，所以这里做了两种格式的兼容（``tools:<id>``
与裸 ``<id>``），并且**取不到 id 时不猜**——宁可把事件归到主 agent，
也不要错标成某个子智能体。
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any, Iterable

from obs import get_logger

logger = get_logger(__name__)

# 子图命名空间里子智能体段的固定前缀（deepagents / langgraph 的约定）
_NS_PREFIX = "tools:"


def parse_namespace(ns: Any) -> str | None:
    """从子图命名空间里取出 task tool_call 的 id。

    ``("tools:abc123",)`` → ``"abc123"``；主图 ``()`` 或无法识别 → ``None``。
    """
    if not ns:
        return None
    if isinstance(ns, str):
        first: Any = ns
    else:
        try:
            first = list(ns)[0]
        except (IndexError, TypeError):
            return None
    if not isinstance(first, str):
        return None
    if first.startswith(_NS_PREFIX):
        candidate = first[len(_NS_PREFIX):]
        return candidate or None
    # 不带前缀的裸 id：兼容命名空间格式变化，但仍要求非空且非已知控制段
    if first and ":" not in first and "/" not in first:
        return first
    return None


@dataclass
class SubagentRun:
    """一次子智能体执行的记录。"""

    call_id: str
    subagent_type: str
    description: str = ""
    status: str = "started"  # started | running | done | error
    started_at: float = field(default_factory=time.time)
    finished_at: float | None = None
    tool_calls: int = 0

    @property
    def duration_ms(self) -> int | None:
        if self.finished_at is None:
            return None
        return int((self.finished_at - self.started_at) * 1000)


class SubagentTracker:
    """一轮请求内的子智能体台账。"""

    def __init__(self) -> None:
        self._runs: dict[str, SubagentRun] = {}
        self._order: list[str] = []

    # -------- 写入 --------

    def note_task_calls(self, tool_calls: Iterable[dict[str, Any]]) -> list[SubagentRun]:
        """从主 agent 的 AIMessage 里登记 ``task`` 工具调用。

        返回本次新登记（以及状态有推进）的子智能体，供上层发事件。
        """
        touched: list[SubagentRun] = []
        for call in tool_calls or []:
            if not isinstance(call, dict):
                continue
            if (call.get("name") or "") != "task":
                continue
            call_id = call.get("id") or ""
            if not call_id:
                # 没有 id 就无法把后续子图事件关联回来，跳过而不是编一个
                logger.debug("task 调用缺少 id，跳过跟踪")
                continue
            args = call.get("args") or {}
            if not isinstance(args, dict):
                args = {}
            subagent_type = str(args.get("subagent_type") or "unknown")
            description = str(args.get("description") or "")

            existing = self._runs.get(call_id)
            if existing is None:
                run = SubagentRun(
                    call_id=call_id,
                    subagent_type=subagent_type,
                    description=description,
                )
                self._runs[call_id] = run
                self._order.append(call_id)
                touched.append(run)
                logger.info(
                    "子智能体启动", extra={"kv": {"event": "subagent.started",
                    "tool_call_id": call_id, "agent": subagent_type}},
                )
            elif existing.status == "started":
                existing.status = "running"
                touched.append(existing)
        return touched

    def note_namespace(self, ns: Any) -> SubagentRun | None:
        """按命名空间标记「该子智能体已在产出」。取不到返回 None。"""
        call_id = parse_namespace(ns)
        if not call_id:
            return None
        run = self._runs.get(call_id)
        if run is None:
            # 表里没有：可能是历史遗留的子图事件（压缩后重建）或框架换了格式。
            # 不新建条目——凭空造一个 subagent_type=unknown 的记录只会污染 UI。
            logger.debug("命名空间 %s 未匹配到已登记的子智能体，事件归主 agent", ns)
            return None
        if run.status == "started":
            run.status = "running"
        run.tool_calls += 1
        return run

    def note_result(self, tool_call_id: str, ok: bool = True) -> SubagentRun | None:
        """登记 task 工具返回，标记子智能体结束。"""
        run = self._runs.get(tool_call_id or "")
        if run is None:
            return None
        run.status = "done" if ok else "error"
        run.finished_at = time.time()
        logger.info(
            "子智能体结束 id=%s type=%s status=%s 耗时=%sms 工具调用=%d",
            run.call_id[:8], run.subagent_type, run.status, run.duration_ms, run.tool_calls,
            extra={"kv": {"event": "subagent.completed", "tool_call_id": run.call_id,
                          "agent": run.subagent_type, "status": run.status, "duration_ms": run.duration_ms}},
        )
        return run

    def finish_all(self, status: str = "done") -> list[SubagentRun]:
        """把所有仍处于活跃状态的子智能体收尾。

        流结束时调用——图可能因为异常或中断提前退出，剩下没收到结果的
        子智能体不能永远挂在 started 状态，否则前端卡片会一直转圈。
        """
        flushed: list[SubagentRun] = []
        for run in self.active():
            run.status = status
            run.finished_at = time.time()
            flushed.append(run)
        return flushed

    # -------- 读取 --------

    def get(self, call_id: str) -> SubagentRun | None:
        """按 task 调用 id 取记录。"""
        return self._runs.get(call_id or "")

    def runs(self) -> list[SubagentRun]:
        """全部记录，按启动顺序。"""
        return [self._runs[cid] for cid in self._order if cid in self._runs]

    def active(self) -> list[SubagentRun]:
        """仍在运行的子智能体。"""
        return [r for r in self.runs() if r.status in ("started", "running")]

    def __len__(self) -> int:
        return len(self._runs)