# -*- coding: utf-8 -*-
"""subagents.py - 后台并行子 agent 委派（2026-09-10，对齐 Hermes delegate_task）

对齐能力：spawn（后台并行，同会话上限 MAX_CONCURRENT=3）/ list / steer（中途追加指示）
/ stop（提前终止，返回部分结果）/ result（取最终结果）。

状态文件化：{workspace}/.sub/{subagent_id}/
  meta.json    {subagent_id, goal, status(running/done/failed/stopped), started_at,
                updated_at(心跳), ended_at, depth, parent_session}
  result.txt   子 agent 的最终输出（流式追加，stop 时即部分结果）
  steer.txt    待消费的追加指示（父 agent 写，子 agent 每轮消费后清空）
  stop.flag    存在即要求停止（子 agent 每轮检测）

文件化的理由与 bg_jobs 相同：web/worker 多进程 + NFS 共享 workspace，
list/steer/stop/result 请求可能落到没跑该子 agent 的进程——文件是唯一跨进程事实源。
steer/stop 通过文件传递，由运行中的子 agent 在下一轮 tool-loop 边界消费。

隔离：
- 目录在发起者 workspace 内（项目工作空间 / _general/{uid}），跨用户/跨项目不可达
- 并发上限按「同一 workspace」计数，机构 A 的子 agent 不占机构 B 的额度
- 子 agent 与父会话历史完全隔离：child config.session_isolation='none'
  （不读父历史、不写 pipeline_conversations——否则子任务对话污染父会话回放）
- 深度限制：子 agent 不能再委派（MAX_DEPTH=1，防递归爆炸烧钱）
- 心跳 stale 检测：meta.updated_at 超 STALE_SECONDS 仍 running = 启动它的
  进程已死，标记 failed（不假装还在跑）
"""
import asyncio
import json
import logging
import os
import re
import time
from dataclasses import replace
from typing import Dict, List, Optional

logger = logging.getLogger("pipeline.subagents")

SUB_DIR = ".sub"
MAX_CONCURRENT = 3          # 同 workspace 并行上限
MAX_DEPTH = 1               # 父=0，子=1；子 agent 禁止再委派
STALE_SECONDS = 300         # 心跳超时判死
RESULT_LIMIT = 20000        # result.txt 回传上限
_ID_RE = re.compile(r"^[A-Za-z0-9_-]{1,32}$")

# 本进程正在跑的子 agent：{subagent_id: asyncio.Task}
_TASKS: Dict[str, asyncio.Task] = {}


class SubagentError(Exception):
    pass


def _validate_id(sid: str) -> str:
    sid = (sid or "").strip()
    if not _ID_RE.match(sid):
        raise SubagentError(f"非法 subagent_id: {sid!r}")
    return sid


def _sub_root(workspace_dir: str) -> str:
    return os.path.join(workspace_dir, SUB_DIR)


def sub_dir(workspace_dir: str, sid: str) -> str:
    return os.path.join(_sub_root(workspace_dir), _validate_id(sid))


def _atomic_write(path: str, text: str):
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        f.write(text)
    os.replace(tmp, path)


def _read_meta(sdir: str) -> dict:
    path = os.path.join(sdir, "meta.json")
    if not os.path.exists(path):
        raise SubagentError("子agent不存在（meta.json 缺失）")
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _write_meta(sdir: str, meta: dict):
    meta["updated_at"] = time.time()
    _atomic_write(os.path.join(sdir, "meta.json"),
                  json.dumps(meta, ensure_ascii=False))


def _stale_guard(meta: dict) -> dict:
    """running 但心跳超时 → 标记 failed（启动进程已死，不假装还在跑）。"""
    if meta.get("status") != "running":
        return meta
    if time.time() - (meta.get("updated_at") or 0) > STALE_SECONDS:
        meta.update({"status": "failed", "ended_at": time.time(),
                     "error": "心跳超时：启动该子agent的进程可能已退出"})
    return meta


# ── 子 agent 运行时钩子（AgentExecutor 每轮调用）──

def heartbeat(workspace_dir: str, sid: str):
    """子 agent 每轮心跳 + 返回待消费指示/停止信号。

    返回 (steer_messages: List[str], stop_requested: bool)。
    文件不存在/解析失败一律当作「无指示、不停止」——钩子绝不能让子 agent 崩。
    """
    steer_msgs: List[str] = []
    stop = False
    try:
        sdir = sub_dir(workspace_dir, sid)
        meta = _read_meta(sdir)
        _write_meta(sdir, meta)            # 心跳
        steer_path = os.path.join(sdir, "steer.txt")
        if os.path.exists(steer_path):
            with open(steer_path, "r", encoding="utf-8") as f:
                raw = f.read().strip()
            if raw:
                steer_msgs = [ln for ln in raw.split("\n@@\n") if ln.strip()]
            os.remove(steer_path)
        if os.path.exists(os.path.join(sdir, "stop.flag")):
            stop = True
    except Exception as e:
        logger.warning(f"subagent heartbeat {sid} failed: {e}")
    return steer_msgs, stop


def append_result(workspace_dir: str, sid: str, text: str):
    """流式追加子 agent 输出（stop 时即为部分结果）。"""
    try:
        path = os.path.join(sub_dir(workspace_dir, sid), "result.txt")
        with open(path, "a", encoding="utf-8") as f:
            f.write(text)
    except Exception as e:
        logger.warning(f"subagent append_result {sid} failed: {e}")


def finish(workspace_dir: str, sid: str, status: str, error: str = ""):
    try:
        sdir = sub_dir(workspace_dir, sid)
        meta = _read_meta(sdir)
        meta.update({"status": status, "ended_at": time.time()})
        if error:
            meta["error"] = error[:500]
        _write_meta(sdir, meta)
    except Exception as e:
        logger.warning(f"subagent finish {sid} failed: {e}")
    _TASKS.pop(sid, None)


# ── 父 agent 侧操作 ──

def count_running(workspace_dir: str) -> int:
    """本 workspace 下 running 子 agent 数（stale 的不计）。"""
    root = _sub_root(workspace_dir)
    if not os.path.isdir(root):
        return 0
    n = 0
    for name in os.listdir(root):
        try:
            meta = _stale_guard(_read_meta(os.path.join(root, name)))
            if meta.get("status") == "running":
                n += 1
        except Exception:
            continue
    return n


async def spawn_subagent(parent_executor, goal: str, context: str = "") -> str:
    """后台派生子 agent，返回 subagent_id。

    parent_executor: 发起委派的 AgentExecutor（继承 user/org/project/pipeline/workspace/config）
    子 agent：独立上下文（不读父历史、不写会话表）、depth+1、禁止再委派。
    """
    goal = (goal or "").strip()
    if not goal:
        raise SubagentError("需要子任务目标 goal")

    depth = getattr(parent_executor, "_delegate_depth", 0)
    if depth >= MAX_DEPTH:
        raise SubagentError("子agent不能再委派子任务（深度限制，防递归爆炸）")

    running = count_running(parent_executor.workspace_dir)
    if running >= MAX_CONCURRENT:
        raise SubagentError(f"并行子agent已达上限 {MAX_CONCURRENT}（当前 {running} 个在跑），"
                            f"先用 subagent(action=list) 查看、等待或 stop 后再派")

    from appPublic.uniqueID import getID
    sid = getID()[:12]
    sdir = sub_dir(parent_executor.workspace_dir, sid)
    os.makedirs(sdir, exist_ok=True)
    _atomic_write(os.path.join(sdir, "result.txt"), "")
    _write_meta(sdir, {
        "subagent_id": sid,
        "goal": goal[:2000],
        "context": (context or "")[:2000],
        "status": "running",
        "depth": depth + 1,
        "started_at": time.time(),
        "ended_at": None,
        "parent_session": getattr(parent_executor, "session_id", "") or "",
        "parent_user": getattr(parent_executor, "user_id", "") or "",
    })

    # 子 executor：同 config 但会话隔离关掉（不读父历史、不写 pipeline_conversations）
    from .agent_loop_v2 import AgentExecutor
    child_config = replace(parent_executor.config, session_isolation="none")
    child = AgentExecutor(
        config=child_config,
        project_id=parent_executor.project_id,
        user_id=parent_executor.user_id,
        workspace_dir=parent_executor.workspace_dir,
        model_name=parent_executor.model_name,
        role=getattr(parent_executor, "role", "") or "",
        session_id="",                      # 不继承父 session（历史/项目上下文都不串）
        generic=getattr(parent_executor, "generic", False),
        default_pipeline_id=getattr(parent_executor, "pipeline_id", "") or "",
        delegate_depth=depth + 1,           # 子agent禁止再委派
    )
    child._subagent_id = sid                # 运行时钩子按此心跳/消费 steer

    prompt = goal if not context else f"{goal}\n\n背景信息：\n{context}"

    async def _runner():
        parts: List[str] = []

        def _flush():
            # 流式落盘：stop/崩溃时 result.txt 即为部分结果（绝不空手）
            _atomic_write(os.path.join(sdir, "result.txt"),
                          "\n".join(x for x in parts if x)[-RESULT_LIMIT:])

        try:
            async for chunk in child.run(prompt):
                try:
                    data = json.loads(chunk)
                except Exception:
                    continue
                t = data.get("type", "")
                if t == "reply":
                    parts.append(data.get("message", ""))
                elif t == "tool_result":
                    # 工具结果摘要（截断），保证中途 stop 也有实质产出可查
                    parts.append("[工具 %s] %s" % (
                        data.get("tool", "?"),
                        (data.get("result", "") or "")[:500]))
                elif t == "error":
                    parts.append(f"[错误] {data.get('message', '')}")
                elif t == "ask_user":
                    # 子 agent 无人可问：记录后终止（父 agent 负责澄清）
                    parts.append(f"[子agent需要澄清，已终止] {data.get('message', '')}")
                    _flush()
                    break
                _flush()
            stopped = os.path.exists(os.path.join(sdir, "stop.flag"))
            if stopped:
                parts.append("[子任务被父agent提前终止，以上为部分结果]")
            _flush()
            finish(parent_executor.workspace_dir, sid,
                   "stopped" if stopped else "done")
        except asyncio.CancelledError:
            parts.append("[子任务被取消，以上为部分结果]")
            try:
                _flush()
                finish(parent_executor.workspace_dir, sid, "stopped")
            finally:
                raise
        except Exception as e:
            logger.exception(f"subagent {sid} crashed")
            parts.append(f"\n[异常] {str(e)[:300]}")
            try:
                _flush()
            finally:
                finish(parent_executor.workspace_dir, sid, "failed", str(e))

    _TASKS[sid] = asyncio.create_task(_runner())
    return sid


def list_subagents(workspace_dir: str) -> List[dict]:
    root = _sub_root(workspace_dir)
    if not os.path.isdir(root):
        return []
    out = []
    for name in sorted(os.listdir(root)):
        try:
            sdir = os.path.join(root, name)
            meta = _stale_guard(_read_meta(sdir))
            if meta.get("status") != _read_meta(sdir).get("status"):
                _write_meta(sdir, meta)      # 落盘 stale 判定
            rpath = os.path.join(sdir, "result.txt")
            size = os.path.getsize(rpath) if os.path.exists(rpath) else 0
            out.append({
                "subagent_id": meta.get("subagent_id", name),
                "goal": (meta.get("goal") or "")[:120],
                "status": meta.get("status"),
                "started_at": meta.get("started_at"),
                "ended_at": meta.get("ended_at"),
                "result_bytes": size,
                "error": meta.get("error", ""),
            })
        except Exception:
            continue
    out.sort(key=lambda d: d.get("started_at") or 0, reverse=True)
    return out


def steer_subagent(workspace_dir: str, sid: str, message: str) -> str:
    message = (message or "").strip()
    if not message:
        raise SubagentError("需要追加指示 message")
    sdir = sub_dir(workspace_dir, sid)
    meta = _stale_guard(_read_meta(sdir))
    if meta.get("status") != "running":
        raise SubagentError(f"子agent已{meta.get('status')}，不能追加指示")
    path = os.path.join(sdir, "steer.txt")
    prev = ""
    if os.path.exists(path):
        with open(path, "r", encoding="utf-8") as f:
            prev = f.read().strip()
    merged = f"{prev}\n@@\n{message}" if prev else message
    _atomic_write(path, merged)
    return "OK: 指示已排队，子agent下一轮消费"


def stop_subagent(workspace_dir: str, sid: str) -> str:
    sdir = sub_dir(workspace_dir, sid)
    meta = _stale_guard(_read_meta(sdir))
    if meta.get("status") != "running":
        return f"子agent已{meta.get('status')}，无需停止"
    _atomic_write(os.path.join(sdir, "stop.flag"), str(time.time()))
    task = _TASKS.get(sid)
    if task is not None and not task.done():
        # 本进程：给 runner 一个取消机会（stop.flag 已由 run 循环优雅处理，
        # 这里只在 LLM 长调用卡住时兜底）
        async def _soft_cancel():
            await asyncio.sleep(90)
            if not task.done():
                task.cancel()
        asyncio.create_task(_soft_cancel())
    return "OK: 已发送停止信号（子agent在当前轮结束后停止，已有部分结果保留）"


def get_result(workspace_dir: str, sid: str) -> dict:
    sdir = sub_dir(workspace_dir, sid)
    meta = _stale_guard(_read_meta(sdir))
    rpath = os.path.join(sdir, "result.txt")
    text = ""
    if os.path.exists(rpath):
        with open(rpath, "r", encoding="utf-8", errors="replace") as f:
            text = f.read()[-RESULT_LIMIT:]
    return {"meta": meta, "result": text or "(暂无输出)"}