"""Harness² 式递归改进，融合进 codex-app-server 的方法层。

来源：Google Research 的 harness2（Apache 2.0，论文 arXiv:2610.14224），
只引它的方法，不引它的部署默认。本模块是按上游思想重写的适配实现，
不是把上游代码原样拷进来：接口按 codex-app-server 的任务/成本模型改过，
提示词是按本项目重写的短版，上游的长提示词、benchmark 适配器、OpenCode /
Claude CLI 通道都没有搬。

一句话：模型权重不动，只改模型周围的 harness（系统提示、技能、脚本、
子代理、钩子），用“提议—探针—合成”循环拿行为证据说话，再把被采纳的
经验写进分领域的经验库，供后面的任务参考。

核心循环（上游叫 Propose-Probe-Compose）：
  1. Propose：对同一任务并行提两类候选 harness——INV（通用做法：流程、
     检查、工具用法）与 DOM（领域做法：该行当的专业知识）。
  2. Probe：每个候选 harness 都在同一任务上真实跑一次（探针），只看
     行为对比，不看任何标准答案——没有标准答案可对，这是开世界任务
     的前提。
  3. Compose：合成器读全部探针轨迹，把有行为证据支撑的改动合成
     一个终版 harness；没有证据的改动不进终版。

两种放量方式（上游的 parallel / sequential）：
  - parallel：一次提 2k 个候选，探针后合成 k 个终版，探得广；
  - sequential：k 轮，每轮提 INV+DOM 各 1 个、合成一个检查点，
    下一轮在检查点上继续改，改得深。

选终版时只比“产物”：候选产物先匿名化（A/B/C……，按任务+k+模式
定种子洗牌，可复现），选择器看不到哪个产物出自哪个 harness；选不
出来时回落到已交付的终版，绝不抛异常把整轮炸掉。k=1 且终版池
只有唯一候选时直接走 trivial：唯一的终版就是答案，不必再选
一次（与选择器是否提供无关）。

经验库（上游叫 harnessing-level recursion）：
  采纳后，反思器只写一条有证据的经验（cue / lesson / evidence 三项
  齐全才入库，缺一项就不存），连同被采纳 harness 的完整树一起存进
  <library>/<domain>/harnesses/<task>/，并重建 INDEX.md。下个同领域
  任务的提议器能读到这些条目——但注意本项目的默认：入库不等于自动
  生效（见下文安全偏差）。

与上游刻意不同的地方（都是为了不破坏本项目已有的安全与计费）：
  1. 不抄上游的执行默认。上游 README 明写：执行器网络不受限，
     Codex 以 sandbox 禁用方式跑。本模块不跑任何命令、不拼任何
     codex exec 参数；执行全部走调用方提供的 adapter，而本项目的
     真实 adapter 必须走 Worker 已有的 Docker / gVisor 沙箱、租户
     隔离与 egress 白名单，绝不加 --dangerously-bypass-*。
  2. 不自动把采纳的 harness 铺到别的任务。codex-app-server 是多租户
     服务，一个编辑器写出的 AGENTS.md 自动改写别人的任务既是质量
     风险也是越权风险。所以 library 默认只是“待评审的经验”：只有
     项目显式开启 auto_apply（后续服务端接线时再加开关，默认关），
     下个任务的 base harness 才会从库里取最新采纳树；默认情况下
     base 永远是项目自己的 base harness，库只给提议器当参考。
  3. 不引 benchmark。任务就是 codex-app-server 里一个普通任务
     （prompt + 产物），没有 JobBench / LAB / WorkBuddy 适配器，
     没有评分器、没有 rubric；delivered 的定义由 adapter 说了算
     （例如“产物文件存在且非空”），选择器只按产物本身判断。
  4. 成本有上限。k 硬上限 3（parallel 一轮最多 1 个 base + 2k 候选
     + k 终版 = 10 次执行；sequential 最多 1 + 3k = 10 次），另有
     可选的 budget_usd：超支后编辑器不再被调用、探针不再执行，
     剩余阶段直接回落到当前 harness，采纳结果也不入库——超支的
     轮次不配进经验库。
  5. 编辑器/选择器/反思器不是本模块内置的模型调用，而是调用方注入
     的三个协议（Editor / Selector / Reflector）。真实接线见
     cloud/harness2_tasks.py：它们各自是一个 codex-app-server 任务，
     prompt 就是本模块生成的 mission，证据以附件 seed 进任务
     worktree（context/ 前缀），编辑器把改动的文件以文件块形式写
     在最终回复里由接线层解析回树；选择器与反思器同理，只解析
     最终回复，不读 worktree 文件（跨主机读不到，见 docs/harness2.md）。

可编辑面（上游 codex substrate 的 surface，原样照搬）：
  - AGENTS.md：常驻文件，每回合都注入，硬上限 8000 字符（上游
    注明 codex 还会在 32768 处静默截断，8000 是留足余量的做法）。
  - .agents/skills/<name>/SKILL.md：技能，frontmatter 必须有
    description，否则加载时会被静默丢弃；且必须在 AGENTS.md 里
    有触发行，否则技能永远不会被调用（本模块会自动补触发行，
    补不下——会超 8000 上限——时跳过，不硬塞）。
  - sub_agents/<name>/AGENT.md：子代理角色，同样需要 description。
  - scripts/*.py：确定性辅助脚本，必须能通过 compile 检查；
    AGENTS.md / 技能里引用的脚本必须真实存在于树里。
  - .codex/hooks.json：生命周期钩子，只接受框架实际派发的事件名，
    且钩子必须是 command 类型、命令非空（格式不对的钩子在框架侧
    是静默不触发，比报错更危险，所以在提议阶段就拦住）。
  不在上述白名单里的路径一律不读、不写、不采纳；想逃出树外的
  路径（..、绝对路径）在规范化时直接拒绝。

用法（确定性演示见 examples/harness2_demo.py，不花模型钱）：
  result = improve_task(adapter, editor, selector, reflector,
                        work_dir=Path("runs/demo"), k=1, mode="parallel")
"""

from __future__ import annotations

import hashlib
import json
import random
import re
import shutil
from dataclasses import dataclass, field
from pathlib import Path
from typing import Protocol, Sequence

# ---------------------------------------------------------------------------
# 可编辑面：codex substrate
# ---------------------------------------------------------------------------

SKILL_DIR = ".agents/skills"
SUBAGENT_DIR = "sub_agents"
SCRIPT_DIR = "scripts"
HOOKS_FILE = ".codex/hooks.json"
ALWAYS_ON = ("AGENTS.md",)
ALWAYS_ON_MAX = 8000
HOOK_EVENTS = frozenset(
    {
        "PreToolUse",
        "PostToolUse",
        "PermissionRequest",
        "PreCompact",
        "PostCompact",
        "SessionStart",
        "SubagentStart",
        "SubagentStop",
        "UserPromptSubmit",
        "Stop",
    }
)

_NAME = r"[a-z0-9][a-z0-9_\-]*"
_SKILL_RX = re.compile(rf"{re.escape(SKILL_DIR)}/({_NAME})/SKILL\.md")
_SKILL_SCRIPT_RX = re.compile(rf"{re.escape(SKILL_DIR)}/{_NAME}/scripts/[A-Za-z0-9_.\-]+")
_SUBAGENT_RX = re.compile(rf"{SUBAGENT_DIR}/({_NAME})/AGENT\.md")
_SCRIPT_RX = re.compile(r"scripts/[a-z0-9_][a-z0-9_]*\.py")
_SCRIPT_REF_RX = re.compile(rf"(?:python3?\s+)((?:{re.escape(SKILL_DIR)}/[\w\-]+/)?scripts/[\w.\-]+\.py)")
_CHOICE_RX = re.compile(r"```choice\s*\n?\s*([A-Za-z0-9_-]+)\s*\n?\s*```", re.S)
_WHY_RX = re.compile(r"<why>(.*?)</why>", re.S)
_NOTES_RX = re.compile(r"<notes>(.*?)</notes>", re.S)
_NO_CHANGE_RX = re.compile(r"<no-change>(.*?)</no-change>", re.S)
_EXP_RX = re.compile(r"<experience>(.*?)</experience>", re.S)
_EDITS_RX = re.compile(r"<edits>(.*?)</edits>", re.S)

Tree = dict[str, str]

# 编辑任务最终回复里的文件块格式（接线层 cloud/harness2_tasks.py 与
# 编辑任务的 prompt 共用这一格式，改一处就要同步另一处）：
#   === HARNESS FILE: <白名单内的相对路径> ===
#   <该文件的完整内容>
#   === END HARNESS FILE ===
# 只输出改动的文件，没改的文件省略即继承（上游同款“省略即继承”
# 契约）。文件内容里若恰好有整行 === END HARNESS FILE === 会被
# 提前截断——极罕见，校验与人工复核是兜底，见 docs/harness2.md。
_HARNESS_FILE_START_RX = re.compile(r"^=== HARNESS FILE:\s*(.+?)\s*===$")
_HARNESS_FILE_END = "=== END HARNESS FILE ==="


def parse_choice(text: str) -> str | None:
    """从选择器最终回复里取 ```choice 块中的匿名 id；没有返回 None。"""
    m = _CHOICE_RX.search(text or "")
    return m.group(1).strip() if m else None


def parse_why(text: str) -> str:
    m = _WHY_RX.search(text or "")
    return m.group(1).strip() if m else ""


def parse_notes(text: str) -> str:
    m = _NOTES_RX.search(text or "")
    return m.group(1).strip() if m else ""


def parse_no_change(text: str) -> str:
    """从编辑任务最终回复里取 <no-change> 标记中的理由；没有返回 ""。

    2026-10-29 两轮真实轮次的教训：编辑在 <notes> 里声称“已修
    正”，却一个文件块都没发，旧契约把这种 notes-only 回复当成
    “无需改动”采纳，导致改动无声丢失。从本版起，无改动必须用
    这个显式标记表态，notes-only 不再算成功（见接线层 edit 的
    格式重试）；标记与文件块同时出现属于自相矛盾，同样算格式错。
    """
    m = _NO_CHANGE_RX.search(text or "")
    return m.group(1).strip() if m else ""


def parse_experience(text: str) -> dict:
    """从反思器最终回复里取 <experience> 与 <edits>。

    cue / lesson / evidence 三项不全，或回复写明 NO_ENTRY 时返回 {}，
    表示没有可入库的经验——调用方不得把半条经验写进库。
    """
    text = text or ""
    if "NO_ENTRY" in text:
        return {}
    m = _EXP_RX.search(text)
    if not m:
        return {}
    out: dict = {}
    for line in m.group(1).strip().splitlines():
        if ":" in line:
            key, _, val = line.partition(":")
            key = key.strip().lower()
            if key in ("cue", "lesson", "evidence", "confidence"):
                out[key] = val.strip()
    if not all(out.get(k) for k in ("cue", "lesson", "evidence")):
        return {}
    edits = _EDITS_RX.search(text)
    if edits:
        out["edits"] = edits.group(1).strip()
    return out


def parse_harness_files(text: str) -> dict[str, str]:
    """解析编辑任务最终回复里的文件块，返回 {相对路径: 完整内容}。

    路径不在可编辑白名单上时直接抛 ValueError——接线层把这当成
    该阶段失败（重试后回落），绝不静默丢掉一块就当没这回事。
    同一路径出现两次时后一块覆盖前一块，与最终回复“最后一次
    表态为准”的直觉一致。
    """
    out: dict[str, str] = {}
    cur_path: str | None = None
    cur_lines: list[str] = []
    for line in (text or "").splitlines():
        m = _HARNESS_FILE_START_RX.match(line.strip())
        if m:
            if cur_path is not None:
                raise ValueError(f"harness file block for {cur_path!r} was not closed")
            cur_path = m.group(1).strip()
            if not is_allowed(cur_path):
                raise ValueError(f"harness file path is off the editable surface: {cur_path!r}")
            cur_lines = []
            continue
        if line.strip() == _HARNESS_FILE_END:
            if cur_path is None:
                raise ValueError("stray === END HARNESS FILE === without a start line")
            out[cur_path] = "\n".join(cur_lines) + ("\n" if cur_lines else "")
            cur_path = None
            cur_lines = []
            continue
        if cur_path is not None:
            cur_lines.append(line)
    if cur_path is not None:
        raise ValueError(f"harness file block for {cur_path!r} was not closed")
    return out


_DELIVERABLE_FILE_START_RX = re.compile(r"^=== DELIVERABLE FILE:\s*(.+?)\s*===$")
_DELIVERABLE_FILE_END = "=== END DELIVERABLE FILE ==="
_DELIVERABLE_FILE_MAX = 8
_DELIVERABLE_FILE_MAX_CHARS = 65536


def _is_safe_deliverable_path(path: str) -> bool:
    if not path or path.startswith("/") or "\x00" in path:
        return False
    return all(p not in ("", ".", "..") for p in path.split("/"))


def parse_deliverable_files(text: str) -> dict[str, str]:
    """解析探针回复里的产物文件块，返回 {相对路径: 内容}。

    与编辑文件块的严格语义不同：这里故意宽容——坏块（未闭合、
    非法路径、超量、超长）只丢弃该块、不抛错，因为产物文件是
    交付文本之外的 enrichment，绝不能因为一块坏内容毁掉整轮
    交付物。超长内容截断并打截断标记，让评审知道它不完整；
    超过 8 块时只取前 8 块。同一路径出现两次时后一块覆盖前一块。
    """
    out: dict[str, str] = {}
    cur_path: str | None = None
    cur_lines: list[str] = []
    for line in (text or "").splitlines():
        m = _DELIVERABLE_FILE_START_RX.match(line.strip())
        if m:
            if cur_path is not None:
                cur_path = None
                cur_lines = []
            path = m.group(1).strip()
            if _is_safe_deliverable_path(path) and len(out) < _DELIVERABLE_FILE_MAX:
                cur_path = path
                cur_lines = []
            continue
        if line.strip() == _DELIVERABLE_FILE_END:
            if cur_path is not None:
                content = "\n".join(cur_lines) + ("\n" if cur_lines else "")
                if len(content) > _DELIVERABLE_FILE_MAX_CHARS:
                    content = (content[:_DELIVERABLE_FILE_MAX_CHARS]
                               + "\n...[truncated]")
                out[cur_path] = content
            cur_path = None
            cur_lines = []
            continue
        if cur_path is not None:
            cur_lines.append(line)
    return out


MAX_K = 3
STAGE_ATTEMPTS = 2  # 提议被拒后给一次按反馈修正的机会，再拒就回落


def normalize_rel(rel: str) -> str:
    """规范化树内相对路径；逃出树外的路径直接拒绝。"""
    r = str(rel).strip().replace("\\", "/")
    if r.startswith("/") or ".." in r.split("/"):
        raise ValueError(f"path escapes the harness: {rel!r}")
    while r.startswith("./"):
        r = r[2:]
    return r.strip("/")


def is_allowed(rel: str) -> bool:
    """该路径是否在可编辑白名单上（永不抛异常，只回答是/否）。"""
    try:
        rel = normalize_rel(rel)
    except ValueError:
        return False
    if rel in ALWAYS_ON or rel == HOOKS_FILE:
        return True
    return bool(
        _SKILL_RX.fullmatch(rel)
        or _SKILL_SCRIPT_RX.fullmatch(rel)
        or _SUBAGENT_RX.fullmatch(rel)
        or _SCRIPT_RX.fullmatch(rel)
    )


def skill_names(tree: Tree) -> list[str]:
    return sorted({m.group(1) for rel in tree if (m := _SKILL_RX.fullmatch(rel))})


def subagent_names(tree: Tree) -> list[str]:
    return sorted({m.group(1) for rel in tree if (m := _SUBAGENT_RX.fullmatch(rel))})


def always_on_text(tree: Tree) -> str:
    return "\n\n".join(tree.get(a, "") for a in ALWAYS_ON)


def read_tree(d: Path | str, *, rejected: list[str] | None = None) -> Tree:
    """读一个目录里的可编辑树；白名单外的文件记入 rejected，不读进来。"""
    d = Path(d)
    out: Tree = {}
    if not d.is_dir():
        return out
    for p in sorted(d.rglob("*")):
        if not p.is_file() or "__pycache__" in p.parts:
            continue
        rel = str(p.relative_to(d))
        if not is_allowed(rel):
            if rejected is not None:
                rejected.append(rel)
            continue
        out[rel] = p.read_text(encoding="utf-8", errors="replace")
    return out


def write_tree(tree: Tree, dest: Path | str) -> Path:
    """把一棵树原样落盘：先清掉目标里旧的可编辑文件，再写新文件。

    目标里白名单外的文件不动——harness 目录里可能还放着任务产物，
    不能因为换 harness 就删掉它们。
    """
    dest = Path(dest)
    dest.mkdir(parents=True, exist_ok=True)
    for p in sorted(dest.rglob("*"), reverse=True):
        if p.is_file() and is_allowed(str(p.relative_to(dest))):
            p.unlink()
    for rel, content in tree.items():
        p = dest / normalize_rel(rel)
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(content, encoding="utf-8")
    return dest


def fingerprint(tree: Tree) -> str:
    """一棵树的摘要：排序后的 路径\\0内容 拼起来算 sha256，前 16 位。"""
    h = hashlib.sha256()
    for rel in sorted(tree):
        h.update(rel.encode())
        h.update(b"\0")
        h.update(tree[rel].encode("utf-8", errors="replace"))
        h.update(b"\0")
    return h.hexdigest()[:16]


def autotrigger(tree: Tree) -> list[str]:
    """给每个技能在 AGENTS.md 里补触发行；会超 8000 上限的行跳过。

    返回补了触发行的技能名。触发行的描述取技能 frontmatter 的
    description 第一行，取不到时用通用措辞。
    """
    fixed: list[str] = []
    for name in skill_names(tree):
        if name in always_on_text(tree):
            continue
        fm = tree.get(f"{SKILL_DIR}/{name}/SKILL.md", "")[:600]
        m = re.search(r"^description:\s*(.+)$", fm, re.M)
        desc = m.group(1).strip().rstrip(".") if m else f"the {name} workflow applies"
        line = f"- When {desc}, read {SKILL_DIR}/{name}/SKILL.md and follow it.\n"
        if len(always_on_text(tree)) + len(line) + 2 > ALWAYS_ON_MAX:
            continue
        cur = tree.get("AGENTS.md", "")
        tree["AGENTS.md"] = (cur.rstrip() + "\n" + line) if cur.strip() else line
        fixed.append(name)
    return fixed


def _check_hooks(src: str, errs: list[str]) -> None:
    try:
        doc = json.loads(src)
    except json.JSONDecodeError as exc:
        errs.append(f"{HOOKS_FILE} is not valid JSON: {exc}")
        return
    hooks = doc.get("hooks") if isinstance(doc, dict) else None
    if not isinstance(hooks, dict) or not hooks:
        errs.append(f'{HOOKS_FILE} needs a top-level {{"hooks": {{"<Event>": [...]}}}} object')
        return
    for event, entries in hooks.items():
        if event not in HOOK_EVENTS:
            errs.append(f"{HOOKS_FILE}: unknown event {event!r}; dispatched events are "
                        f"{', '.join(sorted(HOOK_EVENTS))} (PascalCase)")
            continue
        if not isinstance(entries, list):
            errs.append(f"{HOOKS_FILE}: {event} must map to a LIST of matcher entries")
            continue
        for i, entry in enumerate(entries):
            if not isinstance(entry, dict) or not isinstance(entry.get("hooks"), list):
                errs.append(f"{HOOKS_FILE}: {event}[{i}] needs a `hooks` list")
                continue
            for h in entry["hooks"]:
                if not isinstance(h, dict) or h.get("type") != "command":
                    errs.append(f'{HOOKS_FILE}: {event}[{i}] hook type must be "command"')
                elif not str(h.get("command", "")).strip():
                    errs.append(f"{HOOKS_FILE}: {event}[{i}] has an empty `command`")


def validate(tree: Tree, *, reserved_scripts: frozenset[str] = frozenset()) -> dict:
    """校验一棵树；有问题就抛 ValueError，把全部问题一次列清。

    返回值是树的普查（文件数、技能数、子代理数、脚本数、常驻字符数），
    供日志与测试断言用。校验失败时不返回普查——调用方应按“提议被拒”
    处理，而不是拿半棵树继续跑。
    """
    errs: list[str] = []
    body = always_on_text(tree)
    if len(body) > ALWAYS_ON_MAX:
        errs.append(f"always-on content is {len(body)} chars, over the {ALWAYS_ON_MAX} cap; "
                    "AGENTS.md is injected whole into EVERY turn")
    if not body.strip():
        errs.append("AGENTS.md is empty — it is the only always-on file, "
                    "so an empty one is a harness that does not exist")
    for name in skill_names(tree):
        fm = tree[f"{SKILL_DIR}/{name}/SKILL.md"][:600]
        if not re.search(r"^description:\s*\S", fm, re.M):
            errs.append(f"{SKILL_DIR}/{name}/SKILL.md needs a `description:` in its "
                        "frontmatter, or the skill is dropped SILENTLY at load")
        m_name = re.search(r"^name:\s*(\S+)\s*$", fm, re.M)
        if m_name and m_name.group(1) != name:
            errs.append(f"{SKILL_DIR}/{name}/SKILL.md frontmatter `name: "
                        f"{m_name.group(1)}` OVERRIDES the directory name at load, so "
                        f"every trigger line naming `{name}` stops matching")
        if name not in body:
            errs.append(f"skill '{name}' is never named in the always-on text, so the "
                        "agent will not know to invoke it — an untriggered skill is "
                        "dead weight")
    for name in subagent_names(tree):
        fm = tree[f"{SUBAGENT_DIR}/{name}/AGENT.md"][:600]
        if not re.search(r"^description:\s*\S", fm, re.M):
            errs.append(f"{SUBAGENT_DIR}/{name}/AGENT.md needs a `description:`")
    for rel, src in tree.items():
        if _SCRIPT_RX.fullmatch(rel) or (
            _SKILL_SCRIPT_RX.fullmatch(rel) and rel.endswith(".py")
        ):
            try:
                compile(src, rel, "exec")
            except SyntaxError as exc:
                errs.append(f"{rel} does not compile: {exc.msg} (line {exc.lineno})")
    if HOOKS_FILE in tree:
        _check_hooks(tree[HOOKS_FILE], errs)
    have = {rel for rel in tree if _SCRIPT_RX.fullmatch(rel) or _SKILL_SCRIPT_RX.fullmatch(rel)}
    for rel, src in tree.items():
        if not (_SKILL_RX.fullmatch(rel) or rel in ALWAYS_ON):
            continue
        for ref in _SCRIPT_REF_RX.findall(src):
            if ref.rsplit("/", 1)[-1] in reserved_scripts:
                continue
            if ref not in have:
                errs.append(f"{rel} invokes `{ref}` but that file is not in the harness")
    if errs:
        raise ValueError(" | ".join(errs))
    return {
        "always_on_chars": len(body),
        "n_files": len(tree),
        "skills": len(skill_names(tree)),
        "sub_agents": len(subagent_names(tree)),
        "scripts": sum(1 for rel in tree if _SCRIPT_RX.fullmatch(rel)),
        "hooks": 1 if HOOKS_FILE in tree else 0,
    }


def diff_summary(before: Tree, after: Tree) -> dict:
    return {
        "added": sorted(set(after) - set(before)),
        "removed": sorted(set(before) - set(after)),
        "modified": sorted(rel for rel in set(before) & set(after) if before[rel] != after[rel]),
        "always_on_char_delta": len(always_on_text(after)) - len(always_on_text(before)),
    }


def diversity(trees: dict[str, Tree]) -> dict:
    """候选树之间的两两行 Jaccard 相似度：全塌成同一棵树时能看出来。"""
    names = sorted(trees)
    if len(names) < 2:
        return {"n": len(names), "n_distinct": len(names), "mean_line_jaccard": 1.0}

    def lines(t: Tree) -> set[str]:
        return {f"{rel}:{ln}" for rel, c in t.items() for ln in c.splitlines() if ln.strip()}

    sets = {n: lines(trees[n]) for n in names}
    sims = []
    for i, a in enumerate(names):
        for b in names[i + 1:]:
            u = sets[a] | sets[b]
            sims.append(len(sets[a] & sets[b]) / len(u) if u else 1.0)
    return {
        "n": len(names),
        "n_distinct": len({tuple(sorted(t.items())) for t in trees.values()}),
        "mean_line_jaccard": round(sum(sims) / len(sims), 4),
    }


# ---------------------------------------------------------------------------
# 协议：执行 / 编辑 / 选择 / 反思，全部由调用方注入
# ---------------------------------------------------------------------------


@dataclass
class Rollout:
    label: str
    ok: bool
    delivered: bool
    deliverable_text: str = ""
    trajectory: str = ""
    cost_usd: float = 0.0
    error: str = ""
    deliverable_files: dict[str, str] = field(default_factory=dict)


@dataclass
class EditResult:
    """一次编辑会话的结果。tree 为 None 时，调用方应从 context_dir/harness
    里把编辑器改完的树读回（真实 codex 任务接线就是这个方式）。"""

    tree: Tree | None
    notes: str = ""
    cost_usd: float = 0.0
    cached: bool = False


class TaskAdapter(Protocol):
    """一个任务的执行面：base harness 从哪来、一次探针怎么跑。"""

    task_id: str
    domain: str
    prompt: str

    def base_harness(self) -> Tree:
        """未改动的 base harness。"""
        ...

    def execute(self, label: str, harness: Tree) -> Rollout:
        """用给定 harness 跑一次任务；失败也返回 Rollout，永不抛异常
        （抛了也由 _roll 兜住，记为一次死探针）。"""
        ...

    def validate_tree(self, harness: Tree) -> None:
        """项目特有的额外不变式；默认什么都不拦。"""
        ...


class Editor(Protocol):
    def edit(self, mission: str, context_dir: Path, label: str) -> EditResult:
        """读 context_dir 里的证据（task.md / rollouts / deliverables /
        harness / library），改 context/harness/，返回改完的树。"""
        ...


class Selector(Protocol):
    def select(self, mission: str, context_dir: Path) -> tuple[str, str, float]:
        """返回 (匿名候选 id, 选择理由, 花费 USD)。id 必须是 mission 里
        列出的 A/B/C…… 之一，否则按选择失败回落。"""
        ...


class Reflector(Protocol):
    def reflect(self, mission: str, context_dir: Path) -> dict:
        """返回经验 dict，键至少含 cue / lesson / evidence，可含
        confidence / edits；三项不全时返回 {} 表示没有可入库的经验。"""
        ...


def _adapter_validate(adapter: TaskAdapter, tree: Tree) -> None:
    fn = getattr(adapter, "validate_tree", None)
    if callable(fn):
        fn(tree)


# ---------------------------------------------------------------------------
# 证据上下文：给模型看的文件，而不是把证据塞进 prompt
# ---------------------------------------------------------------------------


def build_context(
    root: Path,
    *,
    task_prompt: str,
    rollouts: Sequence[Rollout],
    harness: Tree,
    library_dir: Path | None = None,
    pool: Sequence[tuple[str, Path, Rollout | None]] = (),
    decisions: dict[str, str] | None = None,
) -> str:
    """在 root 下建 context/ 证据树，返回 mission 里用的布局说明。

    上游的关键做法照搬：证据以文件形式给，prompt 只写指针，这样
    证据不会为了塞进 prompt 而被截断——截断过的轨迹会编造缺陷。
    """
    root = Path(root)
    root.mkdir(parents=True, exist_ok=True)
    (root / "task.md").write_text(task_prompt, encoding="utf-8")
    lines = ["context/task.md                  the task as the agent received it"]
    if rollouts:
        d = root / "rollouts"
        d.mkdir(parents=True, exist_ok=True)
        for r in rollouts:
            (d / f"{r.label}.md").write_text(r.trajectory or "(no trajectory)", encoding="utf-8")
        names = ", ".join(f"{r.label}.md" for r in rollouts[:6])
        more = f" (+{len(rollouts) - 6} more)" if len(rollouts) > 6 else ""
        lines.append(f"context/rollouts/               FULL trajectories: {names}{more}")
        dd = root / "deliverables"
        dd.mkdir(parents=True, exist_ok=True)
        for r in rollouts:
            (dd / f"{r.label}.md").write_text(
                r.deliverable_text or "(no deliverable)", encoding="utf-8"
            )
            for name, content in (r.deliverable_files or {}).items():
                dest = dd / r.label / "files" / name
                dest.parent.mkdir(parents=True, exist_ok=True)
                dest.write_text(content, encoding="utf-8")
        lines.append("context/deliverables/           what each rollout produced"
                     " (+ <label>/files/ when it produced files)")
    write_tree(harness, root / "harness")
    lines.append("context/harness/                the harness tree you may edit, and ONLY this")
    if library_dir and Path(library_dir).is_dir():
        shutil.copytree(library_dir, root / "library", dirs_exist_ok=True)
        lines.append("context/library/                earlier tasks in THIS domain + their harnesses")
    if pool:
        pd = root / "pool"
        pd.mkdir(parents=True, exist_ok=True)
        for label, hdir, rv in pool:
            dest = pd / label
            dest.mkdir(parents=True, exist_ok=True)
            if Path(hdir).is_dir():
                shutil.copytree(hdir, dest / "harness", dirs_exist_ok=True)
            if rv is not None:
                (dest / "rollout.md").write_text(rv.trajectory or "", encoding="utf-8")
                (dest / "deliverable.md").write_text(rv.deliverable_text or "", encoding="utf-8")
        lines.append("context/pool/<label>/           each candidate harness + its rollout")
    if decisions:
        dd = root / "decisions"
        dd.mkdir(parents=True, exist_ok=True)
        for label, note in decisions.items():
            (dd / f"{label}.md").write_text(note or "", encoding="utf-8")
        lines.append("context/decisions/              earlier edits and composition decisions")
    return "\n".join(lines)


def _mission_base(layout: str) -> str:
    return (
        "You are editing an agent harness: the system around a frozen model — "
        "its instructions, skills, scripts, subagents and hooks. You do NOT change "
        "the model, and you do NOT need any ground-truth answer: judge edits by "
        "the behavioral contrasts in the evidence below.\n\n"
        "# EVIDENCE (on disk — read it, do not answer from this message alone)\n\n"
        f"{layout}\n\n"
        "# RULES\n\n"
        "- Edit ONLY files under context/harness/, and only paths on the editable "
        "surface: AGENTS.md, .agents/skills/<name>/SKILL.md, "
        "sub_agents/<name>/AGENT.md, scripts/*.py, .codex/hooks.json. "
        "Anything else you write there is ignored.\n"
        "- AGENTS.md must stay under 8000 chars. Every skill needs a `description:` "
        "in its frontmatter AND a trigger line in AGENTS.md, or it is silently "
        "dropped / never invoked — both are wasted edits.\n"
        "- Scripts must compile, and any script named in AGENTS.md or a skill must "
        "exist in the tree. Hooks must be valid JSON, command-type only, and use "
        "only events the framework dispatches.\n"
        "- Keep what already works. An edit earns its place only when a SPECIFIC "
        "moment in a trajectory shows it changing behavior for the better; "
        "agreement between rollouts is support, not proof.\n"
    )


def _mission_propose(dim: str, layout: str, *, k: int, index: int, round_t: int | None) -> str:
    kind = (
        "INV (invariant): how the agent works in general — process, checks, "
        "tool use, verification habits"
        if dim == "INV"
        else "DOM (domain): what a practitioner in this area knows — domain "
        "facts, conventions, pitfalls"
    )
    where = f"You are candidate {index} of {k} on this dimension." if k > 1 else ""
    round_note = (
        f"\nThis is round {round_t}: you edit the checkpoint the previous round "
        "produced, not the original harness."
        if round_t
        else ""
    )
    return (
        _mission_base(layout)
        + f"\n# YOUR JOB\n\nPropose ONE {kind} edit.\n{where}{round_note}\n"
        "Edit the files in context/harness/ in place, then reply with a "
        "<notes>...</notes> block: what you changed, and the exact trajectory "
        "moment that justifies it. If nothing in the evidence justifies a "
        "change, leave the harness untouched and say so in <notes>."
    )


def _mission_aggregate(layout: str, *, k: int, index: int) -> str:
    return (
        _mission_base(layout)
        + f"\n# YOUR JOB\n\nCompose THE final harness (candidate {index} of {k}). "
        "What you emit is what gets adopted: it runs on this task, and if "
        "selected it enters this domain's library where later tasks read it. "
        "Every rule that bound the candidate proposers binds you — merging an "
        "edit launders it, so only merge edits with trajectory evidence behind "
        "them. Reply with a <notes>...</notes> block listing each merged edit "
        "and its evidence."
    )


def _mission_compose(layout: str, *, t: int, k: int) -> str:
    final = t >= k
    return (
        _mission_base(layout)
        + f"\n# YOUR JOB\n\nCompose checkpoint {t} of {k} in a sequential loop. "
        + (
            "This is the FINAL checkpoint: compose from the WHOLE history, "
            "not just this round — including a candidate rejected earlier, "
            "if a later rollout now argues for it."
            if final
            else "This checkpoint is what the next round will build on."
        )
        + "\nReply with a <notes>...</notes> block listing each merged edit "
        "and its evidence."
    )


def _mission_select(n: int, ids: str) -> str:
    return (
        "You are choosing which of "
        f"{n} candidate outputs should be delivered to the person who asked "
        "for this work. Judge each candidate on what it IS, against the task "
        "and its source material — never on how impressive it sounds, and "
        "never on how it was produced (you are not told that).\n\n"
        "# EVIDENCE\n\n"
        "context/task.md                 the instructions\n"
        "context/candidates/<id>.md      one file per candidate, COMPLETE, "
        "nothing trimmed\n\n"
        f"Open every candidate before choosing. Choose exactly one of: {ids}.\n"
        "Reply with a ```choice block containing only the id, and a "
        "<why>...</why> block explaining the decision in one paragraph."
    )


def _mission_experience() -> str:
    return (
        "You are reading the complete record of one improvement round: a base "
        "harness, the edits proposed against it, the executions those edited "
        "harnesses produced, and which output was finally submitted. Write two "
        "things for a library that future work in this area reads before it "
        "makes its own edits.\n\n"
        "1. <edits>...</edits>: one row per component any candidate changed — "
        "the component, what actually changed, the observed difference in "
        "execution, and a verdict (helped / hurt / no observed effect / "
        "confounded). The adopted candidate is NOT the only evidence: "
        "selection compared deliverables with harnesses hidden, so it says "
        "which output read best, not which edit worked.\n"
        "2. <experience>...</experience>: exactly one transferable lesson, "
        "as `cue: ...`, `lesson: ...`, `evidence: ...`, `confidence: ...` "
        "lines. The lesson must be supported by an observed contrast, not "
        "by a general principle. If no such lesson exists, reply NO_ENTRY "
        "instead of inventing one."
    )


# ---------------------------------------------------------------------------
# 预算：超支后不再花钱，只回落
# ---------------------------------------------------------------------------


@dataclass
class Budget:
    """美元门 + 会话数门（2026-10-29 对照补齐）。

    上游用会话数记账并强校验期望计数（R=4k、P=3k）；本项目此前
    只有美元门，在 unpriced 模型上成本恒为 0、门永远不触发。
    sessions 计探针会话与阶段尝试会话；edit() 内部的 fmt2 纠正
    与选择/反思会话不计入（前者在接线层、后者不走 budget），
    上限语义因此是“受门控的会话数”，不是全轮任务总数。
    """

    max_cost_usd: float | None = None
    spent_usd: float = 0.0
    max_sessions: int | None = None
    sessions: int = 0

    @property
    def exceeded(self) -> bool:
        if self.max_cost_usd is not None and self.spent_usd > self.max_cost_usd:
            return True
        return self.max_sessions is not None and self.sessions >= self.max_sessions

    def add(self, cost: float) -> None:
        self.spent_usd += float(cost or 0.0)

    def add_session(self) -> None:
        self.sessions += 1


# ---------------------------------------------------------------------------
# 阶段：一次提议（含一次修正机会）+ 一次探针
# ---------------------------------------------------------------------------


@dataclass
class Stage:
    label: str
    dim: str
    tree: Tree
    kept_base: bool = False
    fell_back: bool = False
    error: str = ""
    usd: float = 0.0
    attempts: int = 0
    notes: str = ""


@dataclass
class ModeResult:
    mode: str
    task_id: str
    k: int
    stages: list[Stage] = field(default_factory=list)
    trees: dict[str, Tree] = field(default_factory=dict)
    rollouts: dict[str, Rollout] = field(default_factory=dict)
    finals: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)
    total_cost_usd: float = 0.0

    @property
    def ok(self) -> bool:
        return bool(self.finals) and not self.errors

    @property
    def fell_back(self) -> list[str]:
        return [s.label for s in self.stages if s.fell_back]

    def to_dict(self) -> dict:
        return {
            "mode": self.mode,
            "task_id": self.task_id,
            "k": self.k,
            "finals": self.finals,
            "ok": self.ok,
            "fell_back": self.fell_back,
            "kept_base": [s.label for s in self.stages if s.kept_base],
            "total_cost_usd": round(self.total_cost_usd, 6),
            "stages": [
                {
                    "label": s.label,
                    "dim": s.dim,
                    "fell_back": s.fell_back,
                    "kept_base": s.kept_base,
                    "attempts": s.attempts,
                    "usd": s.usd,
                    "notes": s.notes,
                    "error": (s.error or "")[:400],
                }
                for s in self.stages
            ],
            "similarity": diversity({s.label: s.tree for s in self.stages if s.dim == "AGG"})
            if any(s.dim == "AGG" for s in self.stages)
            else None,
            "errors": self.errors,
            # 探针结果此前不在 full_result 里：real5 的终版探针没建
            # 出来时，本地日志完全看不到原因，只能去服务端翻任务。
            # 只放摘要（不放 trajectory），避免日志爆炸。
            "rollouts": {
                lbl: {
                    "ok": r.ok,
                    "delivered": r.delivered,
                    "error": (r.error or "")[:400],
                    "deliverable_text": (r.deliverable_text or "")[:500],
                    "deliverable_files": sorted(r.deliverable_files or {}),
                }
                for lbl, r in self.rollouts.items()
            },
        }


def _accept_tree(adapter: TaskAdapter, new: Tree, fallback: Tree) -> tuple[bool, str]:
    """先自动补触发行，再校验；返回 (是否等于 fallback, 错误信息)。"""
    autotrigger(new)
    kept = not diff_summary(fallback, new)["added"] \
        and not diff_summary(fallback, new)["removed"] \
        and not diff_summary(fallback, new)["modified"]
    try:
        validate(new)
        _adapter_validate(adapter, new)
    except ValueError as exc:
        return kept, str(exc)
    return kept, ""


def _run_stage(
    *,
    label: str,
    dim: str,
    mission: str,
    fallback: Tree,
    adapter: TaskAdapter,
    editor: Editor,
    ctx_dir: Path,
    task_id: str,
    base_rollouts: Sequence[Rollout],
    harness_dir: Path | None,
    library_dir: Path | None,
    pool: Sequence[tuple[str, Path, Rollout | None]] = (),
    decisions: dict[str, str] | None = None,
    budget: Budget | None = None,
) -> Stage:
    if budget is not None and budget.exceeded:
        return Stage(label=label, dim=dim, tree=dict(fallback), fell_back=True,
                     error="budget exceeded before this stage")
    ctx_root = ctx_dir / label
    try:
        layout = build_context(
            ctx_root,
            task_prompt=adapter.prompt,
            rollouts=list(base_rollouts),
            harness=read_tree(harness_dir) if harness_dir else dict(fallback),
            library_dir=library_dir,
            pool=pool,
            decisions=decisions,
        )
    except Exception as exc:
        return Stage(label=label, dim=dim, tree=dict(fallback), fell_back=True,
                     error=f"could not build evidence: {type(exc).__name__}: {exc}"[:400])
    full_mission = (mission.replace("{LAYOUT}", layout) if "{LAYOUT}" in mission
                    else mission + "\n\n# EVIDENCE LAYOUT\n\n" + layout)
    err = ""
    spent = 0.0
    for attempt in range(1, STAGE_ATTEMPTS + 1):
        if budget is not None:
            if budget.exceeded:
                err = "budget exceeded before this attempt"
                break
            budget.add_session()
        try:
            result = editor.edit(
                full_mission + (
                    f"\n\n# YOUR PREVIOUS ATTEMPT WAS REJECTED\n\n{err}\n\n"
                    "Fix it in `context/harness/` and reply again." if err else ""
                ),
                ctx_root,
                f"{label}.a{attempt}",
            )
        except Exception as exc:
            err = f"{type(exc).__name__}: {exc}"
            break
        spent += result.cost_usd
        if budget is not None:
            budget.add(result.cost_usd)
        tree = result.tree
        if tree is None:
            tree = read_tree(ctx_root / "harness")
        if not tree:
            err = "`context/harness/` came back empty — edit the files, do not delete them all"
            continue
        rejected: list[str] = []
        # 编辑器交回的树里若混进了白名单外的文件，只拦不采纳，并告诉它。
        clean = {rel: c for rel, c in tree.items() if is_allowed(rel)}
        if len(clean) != len(tree):
            rejected = sorted(set(tree) - set(clean))
        if rejected:
            err = ("these files are off the editable surface and were not applied: "
                   + ", ".join(rejected[:6]))
            continue
        kept, err = _accept_tree(adapter, clean, fallback)
        if err:
            continue
        notes = result.notes
        return Stage(label=label, dim=dim, tree=clean, kept_base=kept, usd=spent,
                     attempts=attempt, notes=notes)
    return Stage(label=label, dim=dim, tree=dict(fallback), fell_back=True,
                 error=err[:400], usd=spent, attempts=STAGE_ATTEMPTS)


def _roll(
    adapter: TaskAdapter, label: str, tree: Tree, *, budget: Budget | None = None
) -> Rollout:
    if budget is not None and budget.exceeded:
        return Rollout(label=label, ok=False, delivered=False,
                       error="budget exceeded before this rollout")
    if budget is not None:
        budget.add_session()
    try:
        rv = adapter.execute(label, tree)
    except Exception as exc:
        return Rollout(label=label, ok=False, delivered=False,
                       error=f"{type(exc).__name__}: {exc}")
    if budget is not None:
        budget.add(rv.cost_usd)
    return rv


def _roll_many(
    adapter: TaskAdapter,
    items: Sequence[tuple[str, Tree]],
    *,
    budget: Budget | None = None,
) -> dict[str, Rollout]:
    # 探针之间没有共享状态，串行即可；并行是调用方在 adapter 层做的事
    # （例如把每个探针发成一个独立的 codex 任务），不在核心循环里开线程
    # 去抢同一个 Worker/DB——上游的 ThreadPool 在这里会帮倒忙。
    return {label: _roll(adapter, label, tree, budget=budget) for label, tree in items}


def run_parallel(
    task_id: str,
    k: int,
    base_rollouts: Sequence[Rollout],
    *,
    adapter: TaskAdapter,
    editor: Editor,
    work_dir: Path,
    base_tree: Tree,
    library_dir: Path | None = None,
    budget: Budget | None = None,
) -> ModeResult:
    """parallel：2k 个候选探针 + k 个合成终版探针。"""
    result = ModeResult(mode="parallel", task_id=task_id, k=k)
    result.rollouts.update({r.label: r for r in base_rollouts})
    ctx_dir = Path(work_dir) / "context"
    base_dir = write_tree(base_tree, Path(work_dir) / "harnesses" / "base")
    specs = [(f"{d.lower()}{i}", d, i) for i in range(1, k + 1) for d in ("INV", "DOM")]
    stages: list[Stage] = []
    for label, dim, index in specs:
        stages.append(_run_stage(
            label=label, dim=dim,
            mission=_mission_propose(dim, "{LAYOUT}", k=k, index=index, round_t=None),
            fallback=base_tree, adapter=adapter, editor=editor, ctx_dir=ctx_dir,
            task_id=task_id, base_rollouts=base_rollouts, harness_dir=base_dir,
            library_dir=library_dir, budget=budget,
        ))
    # mission 里带 {LAYOUT} 占位：_run_stage 建好上下文后把真实布局
    # 换进去，提议器读到的始终是真实布局，不是空串。
    for stage in stages:
        result.stages.append(stage)
        result.trees[stage.label] = stage.tree
    result.rollouts.update(_roll_many(
        adapter, [(s.label, s.tree) for s in stages], budget=budget))
    pool_ev = [
        (s.label, Path(work_dir) / "harnesses" / s.label, result.rollouts.get(s.label))
        for s in stages
    ]
    # 把候选树落盘，合成器才能在 pool 里读到它们。
    for s in stages:
        write_tree(s.tree, Path(work_dir) / "harnesses" / s.label)
    aggs: list[Stage] = []
    for i in range(1, k + 1):
        aggs.append(_run_stage(
            label=f"agg{i}", dim="AGG",
            mission=_mission_aggregate("{LAYOUT}", k=k, index=i),
            fallback=base_tree, adapter=adapter, editor=editor, ctx_dir=ctx_dir,
            task_id=task_id, base_rollouts=base_rollouts, harness_dir=base_dir,
            library_dir=library_dir, pool=pool_ev,
            decisions={s.label: s.notes or s.error for s in stages}, budget=budget,
        ))
    for stage in aggs:
        result.stages.append(stage)
        result.trees[stage.label] = stage.tree
        result.finals.append(stage.label)
        write_tree(stage.tree, Path(work_dir) / "harnesses" / stage.label)
    result.rollouts.update(_roll_many(
        adapter, [(s.label, s.tree) for s in aggs], budget=budget))
    result.total_cost_usd = sum(s.usd for s in result.stages) + sum(
        r.cost_usd for r in result.rollouts.values()
        if r.label not in {b.label for b in base_rollouts})
    return result


def run_sequential(
    task_id: str,
    k: int,
    base_rollouts: Sequence[Rollout],
    *,
    adapter: TaskAdapter,
    editor: Editor,
    work_dir: Path,
    base_tree: Tree,
    library_dir: Path | None = None,
    budget: Budget | None = None,
) -> ModeResult:
    """sequential：k 轮，每轮 INV+DOM 探针后合成一个检查点。"""
    result = ModeResult(mode="sequential", task_id=task_id, k=k)
    result.rollouts.update({r.label: r for r in base_rollouts})
    ctx_dir = Path(work_dir) / "context"
    current = dict(base_tree)
    history: list[tuple[str, Path, Rollout | None]] = []
    for t in range(1, k + 1):
        current_label = f"c{t - 1}" if t > 1 else "base"
        current_dir = write_tree(current, Path(work_dir) / "harnesses" / current_label)
        stages: list[Stage] = []
        for dim in ("INV", "DOM"):
            stages.append(_run_stage(
                label=f"{dim.lower()}{t}", dim=dim,
                mission=_mission_propose(dim, "{LAYOUT}", k=k, index=t, round_t=t),
                fallback=current, adapter=adapter, editor=editor, ctx_dir=ctx_dir,
                task_id=task_id, base_rollouts=base_rollouts, harness_dir=current_dir,
                library_dir=library_dir, pool=list(history),
                decisions={s.label: s.notes or s.error for s in result.stages},
                budget=budget,
            ))
        for stage in stages:
            result.stages.append(stage)
            result.trees[stage.label] = stage.tree
            write_tree(stage.tree, Path(work_dir) / "harnesses" / stage.label)
        result.rollouts.update(_roll_many(
            adapter, [(s.label, s.tree) for s in stages], budget=budget))
        round_pool = list(history) + [
            (s.label, Path(work_dir) / "harnesses" / s.label, result.rollouts.get(s.label))
            for s in stages
        ]
        ck = f"c{t}"
        comp = _run_stage(
            label=ck, dim="COMPOSE",
            mission=_mission_compose("{LAYOUT}", t=t, k=k),
            fallback=current, adapter=adapter, editor=editor, ctx_dir=ctx_dir,
            task_id=task_id, base_rollouts=base_rollouts, harness_dir=current_dir,
            library_dir=library_dir, pool=round_pool,
            decisions={s.label: s.notes or s.error for s in result.stages}, budget=budget,
        )
        result.stages.append(comp)
        result.trees[ck] = comp.tree
        write_tree(comp.tree, Path(work_dir) / "harnesses" / ck)
        result.rollouts[ck] = _roll(adapter, ck, comp.tree, budget=budget)
        history = round_pool + [(ck, Path(work_dir) / "harnesses" / ck, result.rollouts[ck])]
        current = comp.tree
    result.finals = [f"c{i}" for i in range(1, k + 1)]
    result.total_cost_usd = sum(s.usd for s in result.stages) + sum(
        r.cost_usd for r in result.rollouts.values()
        if r.label not in {b.label for b in base_rollouts})
    return result


def run_mode(mode: str, *args, **kwargs) -> ModeResult:
    token = str(mode or "").strip().lower()
    if token in ("parallel", "agg"):
        return run_parallel(*args, **kwargs)
    if token in ("sequential", "seq"):
        return run_sequential(*args, **kwargs)
    raise ValueError(f"unknown mode {mode!r}; expected parallel | sequential")


# ---------------------------------------------------------------------------
# 选择：只比产物，匿名，失败回落
# ---------------------------------------------------------------------------


@dataclass
class Selection:
    task_id: str
    mode: str
    k: int
    chosen_id: str = ""
    chosen_label: str = ""
    chosen_is_base: bool = False
    why: str = ""
    id_map: dict[str, str] = field(default_factory=dict)
    dropped: list[str] = field(default_factory=list)
    n_candidates: int = 0
    ok: bool = False
    fallback: bool = False
    error: str = ""
    usd: float = 0.0
    trivial: bool = False
    experience: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {k: v for k, v in self.__dict__.items()}


def _anon_id(index: int) -> str:
    out = ""
    n = int(index) + 1
    while n:
        n, rem = divmod(n - 1, 26)
        out = chr(ord("A") + rem) + out
    return out


def assign_ids(labels: Sequence[str], *, task_id: str, k: int, mode: str) -> dict[str, str]:
    """匿名 id -> 真实 label：按排序后定种子洗牌，同输入永远同结果。"""
    # 仅确定性洗牌种子，非安全用途（不防碰撞、不做完整性校验）：
    # usedforsecurity=False 后 bandit B324 不再误报，且摘要值不变。
    seed = int(hashlib.sha1(f"{task_id}|{k}|{mode}".encode(),
                            usedforsecurity=False).hexdigest()[:12], 16)
    ordered = sorted(labels)
    random.Random(seed).shuffle(ordered)
    return {_anon_id(i): lbl for i, lbl in enumerate(ordered)}


def select_final(
    task_id: str,
    k: int,
    mode: str,
    rollouts: dict[str, Rollout],
    *,
    adapter: TaskAdapter,
    selector: Selector | None,
    work_dir: Path,
    base_tree: Tree,
) -> Selection:
    sel = Selection(task_id=task_id, mode=mode, k=k)
    finals = [f"agg{i}" for i in range(1, k + 1)] if mode in ("parallel", "agg") \
        else [f"c{i}" for i in range(1, k + 1)]
    # 只选终版：参考探针（base / 候选）不在池子里，选不到“弃权”。
    pool: dict[str, Rollout] = {}
    for lbl in finals:
        rv = rollouts.get(lbl)
        if rv is None:
            sel.dropped.append(lbl)
            continue
        if not (rv.ok and rv.delivered and rv.deliverable_text.strip()):
            sel.dropped.append(lbl)
            continue
        pool[lbl] = rv
    sel.n_candidates = len(pool)
    if not pool:
        sel.error = "no final produced a deliverable"
        return sel
    sel.id_map = assign_ids(list(pool), task_id=task_id, k=k, mode=mode)
    root = Path(work_dir) / "selection" / f"{mode}{k}"
    shutil.rmtree(root, ignore_errors=True)
    root.mkdir(parents=True, exist_ok=True)
    (root / "task.md").write_text(adapter.prompt, encoding="utf-8")
    cand_dir = root / "candidates"
    cand_dir.mkdir(exist_ok=True)
    for anon, lbl in sel.id_map.items():
        text = pool[lbl].deliverable_text
        header = (f"<!-- candidate {anon}: {len(text):,} chars, "
                  f"{text.count(chr(10)) + 1:,} lines. COMPLETE -- nothing trimmed. -->\n")
        (cand_dir / f"{anon}.md").write_text(header + text, encoding="utf-8")
    ids = ", ".join(sorted(sel.id_map))
    if k == 1 and len(pool) == 1:
        # 唯一终版：不必再开一轮选择会话（上游的 trivial_selection，
        # 条件是“终版池 + k=1”，与选择器是否提供无关；此前误写成
        # selector is None 才 trivial，而接线永远传选择器，导致
        # k=1 也白白跑一个选择任务——2026-10-29 对照发现）。
        label = next(iter(pool))
        sel.ok = True
        sel.trivial = True
        sel.chosen_id = "A"
        sel.chosen_label = label
        sel.id_map = {"A": label}
        sel.why = "trivial pool: single final at k=1; delivered without a selector session"
    else:
        try:
            if selector is None:
                raise ValueError("no selector provided for a multi-candidate pool")
            choice, why, cost = selector.select(_mission_select(len(pool), ids), root)
            sel.usd += float(cost or 0.0)
            if choice not in sel.id_map:
                raise ValueError(f"chose {choice!r}, which is not one of {ids}")
            sel.chosen_id = choice
            sel.chosen_label = sel.id_map[choice]
            sel.why = why
            sel.ok = True
        except Exception as exc:
            # 回落：已交付的终版里排第一的那个；连这个都没有才算选失败。
            sel.chosen_label = sorted(pool)[0]
            sel.chosen_id = next(a for a, lab in sel.id_map.items()
                                 if lab == sel.chosen_label)
            sel.fallback = True
            sel.error = str(exc)[:400]
            sel.ok = True
    sel.chosen_is_base = sel.chosen_label.startswith("base")
    try:
        rec = sel.to_dict()
        (root / "selection.json").write_text(json.dumps(rec, indent=2, default=str), encoding="utf-8")
    except OSError as exc:
        sel.ok = False
        sel.error = (sel.error + f" | could not persist selection: {exc}")[:400]
    return sel


# ---------------------------------------------------------------------------
# 反思与经验库
# ---------------------------------------------------------------------------


def has_experience(experience: dict | None) -> bool:
    return isinstance(experience, dict) and all(
        isinstance(experience.get(key), str) and experience[key].strip()
        for key in ("cue", "lesson", "evidence")
    )


def reflect_adoption(
    task_id: str,
    sel: Selection,
    rollouts: dict[str, Rollout],
    trees: dict[str, Tree],
    *,
    adapter: TaskAdapter,
    reflector: Reflector | None,
    work_dir: Path,
    base_tree: Tree,
    decisions: dict[str, str] | None = None,
) -> dict:
    """采纳后反思：只有反思器交回 cue/lesson/evidence 三项齐全的经验才算数。"""
    if reflector is None or not sel.ok or sel.chosen_label not in trees:
        return {}
    root = Path(work_dir) / "reflection" / f"{sel.mode}{sel.k}"
    layout = build_context(
        root,
        task_prompt=adapter.prompt,
        rollouts=[r for r in rollouts.values() if r.label != "base"],
        harness=trees[sel.chosen_label],
        pool=[(lbl, Path(work_dir) / "harnesses" / lbl, rollouts.get(lbl))
              for lbl in trees if lbl != "base"],
        decisions=decisions,
    )
    write_tree(base_tree, root / "base_harness")
    mission = _mission_experience() + "\n\n# EVIDENCE LAYOUT\n\n" + layout \
        + f"\n\n# ADOPTION\n\nHarness {sel.chosen_label} was adopted. Selection reasoning:\n{sel.why}"
    try:
        experience = reflector.reflect(mission, root)
    except Exception:
        return {}
    if not has_experience(experience):
        return {}
    sel.experience = experience
    return experience


def _slug(task_id: str) -> str:
    return task_id.replace("/", "__")


def _domain_dir(domain: str) -> str:
    return re.sub(r"\s*[/\\]\s*", "__", domain).strip() or "unknown"


class Library:
    """一个领域的经验库：<root>/<domain>/INDEX.md +
    harnesses/<task>/harness/* + EXPERIENCE.md。"""

    def __init__(self, root: Path, domain: str):
        self.root = Path(root)
        self.domain = domain

    @property
    def dir(self) -> Path:
        return self.root / _domain_dir(self.domain)

    @property
    def index(self) -> Path:
        return self.dir / "INDEX.md"

    def entry_dir(self, task_id: str) -> Path:
        return self.dir / "harnesses" / _slug(task_id)

    def exists(self) -> bool:
        return self.index.is_file()

    def context_dir(self) -> Path | None:
        """提议器能读的库目录；空库（没有任何 EXPERIENCE.md）返回 None。"""
        if self.exists() and any(self.dir.rglob("EXPERIENCE.md")):
            return self.dir
        return None

    def entries(self) -> list[str]:
        d = self.dir / "harnesses"
        return sorted(p.name for p in d.iterdir() if p.is_dir()) if d.is_dir() else []

    def tree_of(self, task_id: str) -> Tree:
        return read_tree(self.entry_dir(task_id) / "harness")

    def add(self, task_id: str, harness: Tree, *, why: str = "", label: str = "",
            experience: dict | None = None) -> Path | None:
        """入库一棵被采纳的树 + 采纳理由 + 经验；经验不全时不入库。"""
        if not has_experience(experience):
            return None
        d = self.entry_dir(task_id)
        shutil.rmtree(d / "harness", ignore_errors=True)
        (d / "harness").mkdir(parents=True, exist_ok=True)
        write_tree(harness, d / "harness")
        exp = experience or {}
        lines = [f"# {task_id}", "", f"- adopted: `{label}`", ""]
        if why.strip():
            lines += ["## Why this was chosen", why.strip(), ""]
        if exp.get("edits"):
            lines += ["## What each edit did", str(exp["edits"]).strip(), ""]
        if exp.get("lesson"):
            lines += ["## Transferable lesson"]
            lines += [f"- **{f}**: {str(exp[f]).strip()}"
                      for f in ("cue", "lesson", "evidence", "confidence") if exp.get(f)]
            lines += [""]
        (d / "EXPERIENCE.md").write_text("\n".join(lines), encoding="utf-8")
        self.rebuild_index()
        return d

    def rebuild_index(self) -> Path:
        rows = [
            f"# Library — {self.domain}",
            "",
            "Each entry is an earlier task in THIS domain whose harness was chosen, stored "
            "with its FULL component tree and why it was picked. Browse this index, then "
            "open the entries that look relevant to the matter in front of you.",
            "",
        ]
        names = self.entries()
        if not names:
            rows.append("_(empty — this is the first task in this domain)_")
        for name in names:
            ed = self.dir / "harnesses" / name
            cue = ""
            exp_path = ed / "EXPERIENCE.md"
            if exp_path.is_file():
                for line in exp_path.read_text(encoding="utf-8", errors="replace").splitlines():
                    if line.startswith("- **cue**:"):
                        cue = line.split(":", 1)[1].strip()
            tree = read_tree(ed / "harness")
            rows += [
                f"## `harnesses/{name}/`",
                f"- harness: `harnesses/{name}/harness/` — "
                f"{len(skill_names(tree))} skills, {len(subagent_names(tree))} sub_agents, "
                f"{sum(1 for r in tree if _SCRIPT_RX.fullmatch(r))} scripts; "
                f"{len(always_on_text(tree)):,} always-on chars",
                f"- why:     `harnesses/{name}/EXPERIENCE.md`",
            ]
            if cue:
                rows.append(f"- cue:     {cue}")
            rows.append("")
        self.dir.mkdir(parents=True, exist_ok=True)
        self.index.write_text("\n".join(rows), encoding="utf-8")
        return self.index


# ---------------------------------------------------------------------------
# 一轮完整改进：base 探针 -> 递归 -> 选择 -> 反思 -> 入库
# ---------------------------------------------------------------------------


@dataclass
class ImprovementResult:
    task_id: str
    mode: str
    k: int
    mode_result: ModeResult
    selection: Selection
    experience: dict = field(default_factory=dict)
    library_entry: Path | None = None
    budget_exceeded: bool = False

    @property
    def ok(self) -> bool:
        return self.selection.ok

    def to_dict(self) -> dict:
        return {
            "task_id": self.task_id,
            "mode": self.mode,
            "k": self.k,
            "ok": self.ok,
            "budget_exceeded": self.budget_exceeded,
            "mode_result": self.mode_result.to_dict(),
            "selection": self.selection.to_dict(),
            "experience": self.experience,
            "library_entry": str(self.library_entry) if self.library_entry else None,
        }


def improve_task(
    adapter: TaskAdapter,
    editor: Editor,
    *,
    selector: Selector | None = None,
    reflector: Reflector | None = None,
    work_dir: Path,
    k: int = 1,
    mode: str = "parallel",
    library: Library | None = None,
    auto_apply: bool = False,
    budget_usd: float | None = None,
    max_sessions: int | None = None,
) -> ImprovementResult:
    """跑一轮完整改进。base 探针先跑一次，作为全部候选共享的证据。

    auto_apply=True 时，下个任务的 base 才从库里取最新采纳树——但那
    是调用方在下一轮传 base_harness() 时自己决定的事，本函数只保证
    本轮的 base 永远是 adapter.base_harness() 交来的那棵，不会在轮
    内偷换。默认 False 时库只作提议参考，不改变任何任务的 base。
    """
    if not 1 <= int(k) <= MAX_K:
        raise ValueError(f"k must be 1..{MAX_K}, got {k!r}")
    work_dir = Path(work_dir)
    budget = Budget(max_cost_usd=budget_usd, max_sessions=max_sessions)
    base_tree = adapter.base_harness()
    # 校验 base 本身：base 若已不合法，整轮的对比就没意义，直接报错
    # 而不是让候选在坏 base 上各跑各的。
    validate(base_tree)
    base_rollout = _roll(adapter, "base", base_tree, budget=budget)
    mode_result = run_mode(
        mode, adapter.task_id, int(k), [base_rollout],
        adapter=adapter, editor=editor, work_dir=work_dir, base_tree=base_tree,
        library_dir=library.context_dir() if library else None, budget=budget,
    )
    mode_result.rollouts["base"] = base_rollout
    budget_exceeded = budget.exceeded
    sel = select_final(
        adapter.task_id, int(k), mode_result.mode, mode_result.rollouts,
        adapter=adapter, selector=selector, work_dir=work_dir, base_tree=base_tree,
    )
    experience: dict = {}
    library_entry = None
    if sel.ok and not budget_exceeded:
        experience = reflect_adoption(
            adapter.task_id, sel, mode_result.rollouts, mode_result.trees,
            adapter=adapter, reflector=reflector, work_dir=work_dir, base_tree=base_tree,
            decisions={s.label: s.notes or s.error for s in mode_result.stages},
        )
        if library is not None and has_experience(experience):
            library_entry = library.add(
                adapter.task_id, mode_result.trees[sel.chosen_label],
                why=sel.why, label=sel.chosen_label, experience=experience)
    # auto_apply 目前只是一个显式开关位：调用方在下一轮构造 adapter
    # 时读它来决定 base_harness() 的来源；本轮内不产生任何副作用。
    # 这里把它记在结果里，让调用方和日志都能对上号。
    _ = auto_apply
    return ImprovementResult(
        task_id=adapter.task_id, mode=mode_result.mode, k=int(k),
        mode_result=mode_result, selection=sel, experience=experience,
        library_entry=library_entry, budget_exceeded=budget_exceeded,
    )


__all__ = [
    "SKILL_DIR", "SUBAGENT_DIR", "SCRIPT_DIR", "HOOKS_FILE", "ALWAYS_ON",
    "ALWAYS_ON_MAX", "HOOK_EVENTS", "MAX_K", "Tree",
    "normalize_rel", "is_allowed", "skill_names", "subagent_names",
    "parse_choice", "parse_why", "parse_notes", "parse_no_change", "parse_experience",
    "parse_harness_files",
    "always_on_text", "read_tree", "write_tree", "fingerprint", "autotrigger",
    "validate", "diff_summary", "diversity",
    "Rollout", "EditResult", "TaskAdapter", "Editor", "Selector", "Reflector",
    "Budget", "Stage", "ModeResult", "run_parallel", "run_sequential", "run_mode",
    "Selection", "assign_ids", "select_final",
    "has_experience", "reflect_adoption", "Library",
    "ImprovementResult", "improve_task",
]