#!/usr/bin/env python3
"""抽取智能体派发桥离线测试（阶段 F 第二部分）。

固定合同：任务包装配（正文注入/白名单/修复提示词）、单轮派发产出合法抽取、
机械校验失败触发一轮修复重派、两轮失败走保守收口、派发失败失败关闭。
"""
from __future__ import annotations

import contextlib
import json
import pathlib
import sys
import tempfile
import unittest
from unittest import mock

PROJECT_ROOT = pathlib.Path(__file__).resolve().parents[3]
SCRIPT_DIR = PROJECT_ROOT / ".agent" / "skills" / "extract-chapter-knowledge" / "scripts"
for path in (SCRIPT_DIR,):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

import dispatch_extraction_bridge as bridge  # noqa: E402
from dispatch_extraction_bridge import (  # noqa: E402
    DispatchExtractionError,
    build_extraction_task_spec,
    extractor_session_paths,
    run_extraction_via_dispatch,
)

BODY = "第三章 波纹。茧从实验舰外舱门的裂口挤出去，金属边缘像被剥开的肋骨。林深盯着深渊方向。"


def _valid_payload(evidence="茧从实验舰外舱门的裂口挤出去"):
    return {
        "entities": [
            {"type": "character", "name": "林深", "brief": "试机师", "fields": {}, "evidence": "林深盯着深渊方向"},
            {"type": "item", "name": "茧", "brief": "生物机甲", "fields": {}, "evidence": evidence},
        ],
        "relations": [
            {"source": "林深", "target": "茧", "type": "神经链接", "description": "链接",
             "evidence": "茧从实验舰外舱门的裂口挤出去"},
        ],
        "state": {"currentSituation": "出击", "characterStates": {}, "foreshadowing": {"埋": [], "推": [], "收": []}, "handoff": ""},
    }


def _load_chapter_fake(work_id, chapter_order):
    return ("深渊机神", 301, "波纹", 901, BODY), ["林深", "茧"]


class ExtractionSpecTest(unittest.TestCase):

    def test_spec_shape(self):
        spec = build_extraction_task_spec(
            work_id=12, chapter_order=3, title="深渊机神", chapter_title="波纹",
            existing_names=["林深"], body=BODY,
        )
        self.assertEqual(spec["role"], "extractor")
        self.assertEqual(spec["input"]["body"], BODY)
        self.assertEqual(spec["input"]["existingEntityNames"], ["林深"])
        self.assertEqual(spec["toolAllowlist"], ["read_chapter_text", "search_entities"])
        self.assertIn("逐字连续片段", spec["taskPrompt"])
        self.assertNotIn("机械校验失败", spec["taskPrompt"])

    def test_spec_repair_appends_reason(self):
        spec = build_extraction_task_spec(
            work_id=12, chapter_order=3, title="深渊机神", chapter_title="波纹",
            existing_names=[], body=BODY, repair_reason="entities[0] 证据不在正文中",
        )
        self.assertIn("机械校验失败", spec["taskPrompt"])
        self.assertIn("不得新增条目", spec["taskPrompt"])

    def test_session_paths(self):
        sid, path = extractor_session_paths(12, 3)
        self.assertEqual(sid, "extractor-work12-ch3")
        self.assertIn("extractor-sessions", str(path))


class OrchestrationTest(unittest.TestCase):

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.tmp = pathlib.Path(self._tmp.name)
        self.dispatches = []

    def tearDown(self):
        self._tmp.cleanup()

    def _fake_dispatch(self, outputs):
        """outputs: 每次派发依次返回的 JSON（Exception 表示派发失败）。"""
        outputs = list(outputs)

        def fake(spec_file, **kwargs):
            self.dispatches.append(kwargs["run_id"])
            run_dir = self.tmp / kwargs["run_id"]
            run_dir.mkdir(parents=True)
            out = outputs[len(self.dispatches) - 1]
            if isinstance(out, Exception):
                return {"status": "failed", "errorCode": "TIMEOUT"}, 1
            (run_dir / "output.json").write_text(json.dumps(out, ensure_ascii=False), encoding="utf-8")
            return {"status": "completed", "runDir": str(run_dir),
                    "actualModelIds": ["catproxy-anthropic/claude-opus-5"],
                    "usage": {"input_tokens": 10}}, 0

        return fake

    @contextlib.contextmanager
    def _patched(self, outputs):
        fake = self._fake_dispatch(outputs)
        with mock.patch.object(bridge, "run_dispatch", fake), \
             mock.patch.object(bridge, "_load_chapter", _load_chapter_fake), \
             mock.patch.object(bridge, "start_run", lambda **kw: {"run_id": kw["run_id"]}), \
             mock.patch.object(bridge, "finish_run", lambda *a, **kw: None), \
             mock.patch.object(bridge, "record_failure", lambda *a, **kw: None), \
             mock.patch.object(bridge, "_propose_chapter_extract_lesson", lambda **kw: None), \
             mock.patch.object(bridge, "persist_extraction",
                               lambda *a, **kw: {"draft_ids": [1, 2, 3], "state_draft_id": 9}):
            yield

    def _run(self, outputs):
        with self._patched(outputs):
            return run_extraction_via_dispatch(
                12, 3, repo_root=self.tmp, provider="catproxy-anthropic",
                model="claude-opus-5", run_id="run-extract-test", spec_dir=self.tmp,
            )

    def test_single_dispatch_success(self):
        summary = self._run([_valid_payload()])
        self.assertEqual(summary["run_id"], "run-extract-test")
        self.assertEqual(len(summary["draft_ids"]), 3)
        self.assertEqual(self.dispatches, ["run-extract-test-extractor-v1"])

    def test_repair_round_when_evidence_invalid(self):
        bad = _valid_payload(evidence="不在正文里的证据")
        summary = self._run([bad, _valid_payload()])
        self.assertEqual(self.dispatches, ["run-extract-test-extractor-v1", "run-extract-test-extractor-v2"])
        self.assertEqual(len(summary["draft_ids"]), 3)

    def test_salvage_after_two_failures(self):
        bad = _valid_payload(evidence="不在正文里的证据")
        summary = self._run([bad, bad])
        # salvage：证据无效但实体名逐字在正文 → 保留并以实体名为最小证据；关系缺证据丢弃
        self.assertEqual(len(summary["draft_ids"]), 3)
        self.assertEqual(self.dispatches, ["run-extract-test-extractor-v1", "run-extract-test-extractor-v2"])

    def test_dispatch_failure_fails_closed(self):
        with self.assertRaises(DispatchExtractionError):
            self._run([TimeoutError()])


if __name__ == "__main__":
    unittest.main()