# -*- coding: utf-8 -*-
"""
orchestrator.py — Pi Agent 兜底能力 P0 PoC：Python 编排器
=========================================================

职责（对应 `docs/architecture/fallback.md` 的 FallbackLane 编排与熔断设计）：

1. 以子进程拉起 Pi headless（`node .../cli.js -p --mode json`），逐行解析 JSONL 事件流；
2. 三重熔断：超时 / 步数上限 / 输出体量上限，触发即杀进程树（Windows taskkill /T /F）
   并把运行显式标记为 failed（绝不把失败包装成成功 —— GOAL.md 价值观）；
3. 环境变量白名单清洗后传给子进程（剥离 CONDA_*/PYTHONHOME/PYTHONPATH 等），
   并设置 PI_CODING_AGENT_DIR 指向 poc 内配置目录（NOTES-pi-runtime.md §3）；
4. 成败判定基于事件流里的 stopReason（"stop"=正常 / "error"=失败），
   **绝不相信进程退出码**（pi headless 失败也退 0 —— Agent-A 实测坑 #1）；
5. 全程写运行日志到 runs/<runId>/：events.jsonl（原始事件流）、orchestrator.log、result.json。

mock 模式说明：AgentRunner 是协议抽象，mock executor（demo_task.py）产出与真实
pi 相同结构的事件序列，因此熔断/日志/stopReason 判定走**同一条代码路径**。
"""

from __future__ import annotations

import json
import os
import queue
import subprocess
import threading
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Iterable, Iterator, Optional

from sandbox import POC_ROOT, RUNS_ROOT, PI_HOME, PI_CLI, create_run_dirs

# ---------------------------------------------------------------------------
# 集中配置（改行为只改这里）
# ---------------------------------------------------------------------------

@dataclass
class BreakerConfig:
    """三重熔断阈值（对应方案 §4.6「熔断三闸」）。"""
    timeout_sec: float = 600.0            # 闸 1：单次运行超时（默认 10min，方案 §S10）
    max_steps: int = 50                   # 闸 2：工具调用步数上限（方案 §4.6 步数 50）
    max_output_bytes: int = 10 * 1024 * 1024  # 闸 3：assistant 输出累计体量上限（10MiB）
    poll_interval_sec: float = 1.0        # 读事件流的轮询间隔（也用于检测进程挂起）


# 子进程环境变量白名单（L4 进程层围墙，对应方案 §4.2 buildChildEnv 语义）
ENV_WHITELIST = (
    "PATH", "PATHEXT", "SYSTEMROOT", "SYSTEMDRIVE", "WINDIR", "COMSPEC",
    "TEMP", "TMP", "USERPROFILE", "APPDATA", "LOCALAPPDATA", "HOME",
    "NODE_OPTIONS",  # node 运行时可能需要；其余一律不传
)

# 显式剥离名单（即使同名出现在白名单也再剥一层，双保险）
ENV_STRIP_PREFIXES = ("CONDA_", "PYTHON", "PIP_", "VIRTUAL_ENV")


def build_child_env(extra: Optional[dict] = None) -> dict:
    """
    构造传给 Pi 子进程的清洗后环境（白名单制）。

    - 只放行 ENV_WHITELIST 里的变量；
    - 剥离 CONDA_*/PYTHONHOME/PYTHONPATH/PIP_*/VIRTUAL_ENV*（防止 Anaconda 污染 node 子进程）；
    - 强制设置 PI_CODING_AGENT_DIR 指向 poc 内 runtime/pi-home（配置圈禁，Agent-A 实测机制）。
    """
    env: dict = {}
    for key in ENV_WHITELIST:
        if key in os.environ and not key.upper().startswith(ENV_STRIP_PREFIXES):
            env[key] = os.environ[key]
    env["PI_CODING_AGENT_DIR"] = str(PI_HOME)
    if extra:
        for k, v in extra.items():
            if not k.upper().startswith(ENV_STRIP_PREFIXES):
                env[k] = v
    return env


# ---------------------------------------------------------------------------
# 运行结果
# ---------------------------------------------------------------------------

@dataclass
class RunResult:
    """一次兜底运行的最终判定。success 只由 stopReason 决定。"""
    run_id: str
    success: bool
    stop_reason: str = ""            # "stop" / "error" / "breaker:<哪种熔断>" / "harness_error"
    error_message: str = ""
    steps: int = 0                   # 工具调用步数（tool_execution_start 计数）
    output_bytes: int = 0            # assistant 输出累计字节
    elapsed_sec: float = 0.0
    final_text: str = ""             # 最后一条 assistant 消息文本（即 Pi 的报告）
    run_dir: str = ""
    extra: dict = field(default_factory=dict)


# ---------------------------------------------------------------------------
# 事件流来源抽象
# ---------------------------------------------------------------------------

# AgentRunner：给定 (task, work_dir, event_sink) 产出 JSONL 事件 dict 迭代器。
# 真实实现 = pi headless 子进程；mock 实现 = demo_task.py 的纯 Python executor。
AgentRunner = Callable[[str, Path], Iterator[dict]]


class Orchestrator:
    """编排器本体：消费事件流，执行熔断与判定，落日志。"""

    def __init__(self, config: Optional[BreakerConfig] = None):
        self.cfg = config or BreakerConfig()

    # -- 主入口 -----------------------------------------------------------

    def run(self, task: str, runner: AgentRunner, run_id: Optional[str] = None,
            on_tool_event: Optional[Callable[[dict], None]] = None) -> RunResult:
        """
        执行一次兜底运行。

        参数:
            task:         任务简报（自然语言）。
            runner:       AgentRunner，产出事件流（真实或 mock 共用本条熔断路径）。
            run_id:       缺省自动生成；目录 runs/<run_id>/ 由 sandbox.create_run_dirs 建好。
            on_tool_event: 每个 tool_execution_start/end 事件的回调（tool_bridge 借此签发 callId）。

        返回 RunResult；任何熔断/错误都会显式写进 result.json 与 orchestrator.log。
        """
        run_id = run_id or time.strftime("run-%Y%m%d-%H%M%S-") + uuid.uuid4().hex[:6]
        dirs = create_run_dirs(run_id)
        run_dir = dirs["root"]
        log_path = run_dir / "orchestrator.log"
        events_path = run_dir / "events.jsonl"

        def log(msg: str) -> None:
            line = f"[{time.strftime('%H:%M:%S')}] {msg}\n"
            with open(log_path, "a", encoding="utf-8") as f:
                f.write(line)

        res = RunResult(run_id=run_id, success=False, run_dir=str(run_dir))
        t0 = time.monotonic()
        log(f"run_id={run_id} task={task[:120]!r}")
        log(f"breaker config: {self.cfg}")

        breaker_tripped: Optional[str] = None
        last_assistant_text = ""
        stop_reason = ""
        error_message = ""

        try:
            with open(events_path, "w", encoding="utf-8") as evf:
                for event in runner(task, dirs["work"]):
                    elapsed = time.monotonic() - t0

                    # —— 闸 1：超时 ——
                    if elapsed > self.cfg.timeout_sec:
                        breaker_tripped = f"breaker:timeout({elapsed:.1f}s>{self.cfg.timeout_sec}s)"
                        self._safe_write_event(evf, {"type": "breaker", "reason": breaker_tripped})
                        break

                    self._safe_write_event(evf, event)
                    etype = event.get("type", "")

                    # —— 步数统计 + 闸 2 ——
                    if etype == "tool_execution_start":
                        res.steps += 1
                        if on_tool_event:
                            on_tool_event(event)
                        if res.steps > self.cfg.max_steps:
                            breaker_tripped = f"breaker:max_steps({res.steps}>{self.cfg.max_steps})"
                            self._safe_write_event(evf, {"type": "breaker", "reason": breaker_tripped})
                            break
                    elif etype == "tool_execution_end" and on_tool_event:
                        on_tool_event(event)

                    # —— 输出体量统计 + 闸 3 ——
                    text_delta = self._extract_text_delta(event)
                    if text_delta:
                        res.output_bytes += len(text_delta.encode("utf-8"))
                        if res.output_bytes > self.cfg.max_output_bytes:
                            breaker_tripped = (
                                f"breaker:max_output({res.output_bytes}>{self.cfg.max_output_bytes})"
                            )
                            self._safe_write_event(evf, {"type": "breaker", "reason": breaker_tripped})
                            break

                    # —— stopReason 判定（成败唯一权威，Agent-A 坑 #1）——
                    sr, em, txt = self._extract_stop(event)
                    if sr:
                        stop_reason, error_message = sr, em
                        if txt:
                            last_assistant_text = txt

                    # 自动重试事件：只记录（Agent-A 坑 #4，重试 ~14s 已计入超时预算）
                    if etype == "auto_retry_start":
                        log(f"auto_retry_start attempt={event.get('attempt')}")

                    if etype == "agent_end":
                        break

        except Exception as exc:  # runner 自身抛错（如子进程 spawn 失败）
            stop_reason = "harness_error"
            error_message = f"{type(exc).__name__}: {exc}"
            log(f"HARNESS ERROR: {error_message}")

        res.elapsed_sec = time.monotonic() - t0
        res.final_text = last_assistant_text

        if breaker_tripped:
            res.stop_reason = breaker_tripped
            res.error_message = "熔断触发，运行显式标记失败"
            log(f"BREAKER TRIPPED: {breaker_tripped} -> failed")
        elif stop_reason == "stop":
            res.success = True
            res.stop_reason = "stop"
            log(f"OK stopReason=stop steps={res.steps} out={res.output_bytes}B "
                f"elapsed={res.elapsed_sec:.1f}s")
        else:
            res.stop_reason = stop_reason or "error:no_stop_reason"
            res.error_message = error_message or "事件流未给出 stopReason=stop，按失败处理"
            log(f"FAILED stopReason={res.stop_reason} err={res.error_message}")

        with open(run_dir / "result.json", "w", encoding="utf-8") as f:
            json.dump(res.__dict__, f, ensure_ascii=False, indent=2)
        return res

    # -- 事件解析辅助 ------------------------------------------------------

    @staticmethod
    def _safe_write_event(evf, event: dict) -> None:
        try:
            evf.write(json.dumps(event, ensure_ascii=False) + "\n")
        except Exception:
            evf.write(json.dumps({"type": "unserializable_event"}) + "\n")

    @staticmethod
    def _extract_text_delta(event: dict) -> str:
        """从 message_update/message_end 事件取 assistant 文本增量（用于输出体量闸）。"""
        if event.get("type") == "message_update":
            delta = event.get("delta") or {}
            if isinstance(delta, dict):
                return str(delta.get("text") or "")
        if event.get("type") == "message_end":
            msg = event.get("message") or {}
            for part in (msg.get("content") or []):
                if isinstance(part, dict) and part.get("type") == "text":
                    return str(part.get("text") or "")
        return ""

    @staticmethod
    def _extract_stop(event: dict):
        """
        从事件里提取 (stopReason, errorMessage, 文本)。
        兼容两个位置：message_end.message.stopReason 与 agent_end.messages 最后一条。
        """
        etype = event.get("type")
        msg = None
        if etype == "message_end":
            msg = event.get("message") or {}
        elif etype == "agent_end":
            msgs = event.get("messages") or []
            assistants = [m for m in msgs if m.get("role") == "assistant"]
            msg = assistants[-1] if assistants else None
        if not msg:
            return None, None, None
        txt = ""
        for part in (msg.get("content") or []):
            if isinstance(part, dict) and part.get("type") == "text":
                txt += str(part.get("text") or "")
        return msg.get("stopReason"), msg.get("errorMessage"), txt

    # -- 真实 pi headless runner ------------------------------------------

    def pi_headless_runner(self, model: str = "aps-local/qwen3-32b-local",
                           tools: str = "read,grep,find,ls",
                           guard_ext: Optional[Path] = None,
                           extra_args: Optional[list] = None) -> AgentRunner:
        """
        构造真实 pi headless 的 AgentRunner（供 --real 模式，Agent-C 使用）。

        命令形态（Agent-A 实测）：
            node <PI_CLI> -p --mode json --model <m> --tools <白名单> [-e guard.ts] "task"
        - cwd 圈禁在 runs/<runId>/work（L2 文件层，pi 自身无沙箱）；
        - 环境经 build_child_env 清洗（L4）；
        - 子进程 stdout 由读线程转 queue，主循环带 poll 超时 —— 进程挂起也能被闸 1 抓到；
        - 熔断/结束时 taskkill /T /F 杀整棵进程树。
        """
        cmd_tail = extra_args or []

        def runner(task: str, work_dir: Path) -> Iterator[dict]:
            cmd = [
                "node", str(PI_CLI),
                "-p", "--mode", "json",
                "--model", model,
                "--tools", tools,
            ]
            if guard_ext:
                cmd += ["-e", str(guard_ext)]
            cmd += cmd_tail + [task]

            env = build_child_env()
            proc = subprocess.Popen(
                cmd, cwd=str(work_dir), env=env,
                stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
            )
            q: "queue.Queue[Optional[str]]" = queue.Queue()

            def reader():
                try:
                    for raw in proc.stdout:
                        q.put(raw.decode("utf-8", errors="replace"))
                finally:
                    q.put(None)  # EOF 哨兵

            t = threading.Thread(target=reader, daemon=True)
            t.start()

            killed = False
            try:
                while True:
                    try:
                        line = q.get(timeout=self.cfg.poll_interval_sec)
                    except queue.Empty:
                        if proc.poll() is not None and t.is_alive() is False:
                            break
                        # 让上层熔断循环有机会检查超时：产出一个心跳事件
                        yield {"type": "harness_heartbeat"}
                        continue
                    if line is None:
                        break
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        yield json.loads(line)
                    except json.JSONDecodeError:
                        yield {"type": "harness_nonjson", "raw": line[:500]}
            finally:
                # 进程回收（L4）：编排器熔断 break 时会 close 本生成器，
                # 这里只杀进程树、不许再 yield（closing 中的生成器 yield 会 RuntimeError）。
                if proc.poll() is None:
                    killed = True
                    kill_process_tree(proc.pid)
                try:
                    proc.wait(timeout=15)
                except Exception:
                    pass
            # 以下仅正常 EOF 退出时可达（被 close 时 GeneratorExit 会直接传播出去）
            if killed:
                yield {"type": "harness_killed", "pid": proc.pid}
            err = proc.stderr.read().decode("utf-8", errors="replace") if proc.stderr else ""
            if err.strip():
                yield {"type": "harness_stderr", "text": err[-2000:]}

        return runner


def kill_process_tree(pid: int) -> None:
    """Windows 下杀整棵进程树（node 起的孙进程一并回收，对应方案 §4.2 L4 进程回收）。"""
    subprocess.run(
        ["taskkill", "/PID", str(pid), "/T", "/F"],
        capture_output=True, timeout=30,
    )