from __future__ import annotations

from pathlib import Path
import re
import sys


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from swarm_minimal.azure_store import PostgresRedisBlobSwarmStore
from swarm_minimal.config import SwarmConfig
from swarm_minimal.local_env import load_project_env


RUN_IDS = [
    "04c641d170fe4ea7aa3d882d9df37cca",
    "c592a7ca4c0f4e02b93a1390b62d0af7",
    "063632eeb17b45c197aa866066158667",
    "3e8e58ae4e084bc8b90cf5c46f8992f3",
    "78f189ccd1924ed0a4fb0a0a447ad449",
    "623b5f6e5cc24cc7967fd9577f9c224b",
    "42688e7b6245466dab8398dfe4790456",
]

TEST_SCENARIOS = (
    {
        "id": "S01",
        "name": "syntax_import_sanity",
        "purpose": "验证所有 Python 模块可编译，排除语法和导入层错误。",
        "evidence": "`py_compile swarm_minimal/*.py examples/*.py tests/*.py`",
    },
    {
        "id": "S02",
        "name": "unit_regression",
        "purpose": "验证内存蜂群、NewAPI mock、配置脱敏和基础收敛行为。",
        "evidence": "`unittest discover -s tests`",
    },
    {
        "id": "S03-S06",
        "name": "deterministic_standard_scenarios",
        "purpose": "验证链路连续性、依赖边界、最终评分和失败注入。",
        "evidence": "`unittest tests.test_standard_scenarios`",
    },
    {
        "id": "S07",
        "name": "live_azure_newapi_continuous_reasoning",
        "purpose": "用真实 Azure PostgreSQL、Redis、Blob 和 NewAPI 跑 7 步连续推理链。",
        "evidence": "本报告下方每个 live run 的 task.input / task.output / handoff 记录。",
    },
    {
        "id": "S08",
        "name": "model_io_report_audit",
        "purpose": "验证本报告包含场景、输入、输出、交接证据，且没有明显真实密钥样式。",
        "evidence": "`unittest tests.test_model_io_report_audit`",
    },
)

OUTPUT_PATH = ROOT / "docs" / "MODEL_AGNET_IO_REPORT.zh-CN.md"


def main() -> None:
    load_project_env(ROOT)
    store = PostgresRedisBlobSwarmStore(SwarmConfig.from_env())
    try:
        store.ensure_schema()
        report = build_report(store)
        OUTPUT_PATH.parent.mkdir(exist_ok=True)
        OUTPUT_PATH.write_text(report, encoding="utf-8")
        print(str(OUTPUT_PATH))
    finally:
        store.close()


def build_report(store: PostgresRedisBlobSwarmStore) -> str:
    sections = [
        "# 模型 / Agnet 输入输出报告",
        "",
        "这份报告直接从 Azure PostgreSQL 的 `swarm_convergence`、`swarm_tasks` 和 `swarm_observations` 读取历史 live run。",
        "报告只展示任务输入、模型提示词模板、Agnet 输出和评分，不展示 `.env` 或任何密钥。",
        "",
        "## 本轮测试场景补充",
        "",
        "本报告重点解释 S07 live 场景里的模型输入输出，但它属于完整标准矩阵的一部分；本轮重跑覆盖以下场景：",
        "",
        "| ID | 场景 | 测试目的 | 证据入口 |",
        "| --- | --- | --- | --- |",
    ]
    for scenario in TEST_SCENARIOS:
        sections.append(
            f"| {scenario['id']} | {scenario['name']} | {scenario['purpose']} | {scenario['evidence']} |"
        )
    sections.extend(
        [
            "",
            "最新重跑结论：S01-S08 全部 PASS；最新 S07 live run 会排在下方第一个。",
            "",
        ]
    )
    for run_id in RUN_IDS:
        convergence = fetch_convergence(store, run_id)
        if convergence is None:
            sections.extend([f"## Run `{run_id}`", "", "未找到该 run。", ""])
            continue
        task_ids = [item["task_id"] for item in convergence["observations"]]
        tasks = fetch_tasks(store, task_ids)
        sections.extend(render_run(convergence, tasks))
    return "\n".join(line.rstrip() for line in "\n".join(sections).splitlines()).rstrip() + "\n"


def render_run(convergence: dict[str, object], tasks: dict[str, dict[str, object]]) -> list[str]:
    run_id = str(convergence["run_id"])
    goal = str(convergence["goal"])
    system_prompt = system_prompt_for_goal(goal)
    lines = [
        f"## Run `{run_id}`",
        "",
        f"- 任务目标：{goal}",
        f"- 完成任务数：{convergence['completed_tasks']}",
        f"- 收敛分数：{convergence['accepted_score']}",
        f"- Blob artifact：`{convergence['artifact_path']}`",
        "",
        "### 模型系统提示词",
        "",
        "```text",
        system_prompt,
        "```",
        "",
    ]

    for index, observation in enumerate(convergence["observations"], start=1):
        task_id = observation["task_id"]
        task = tasks[task_id]
        output = redact_sensitive_text(str(task.get("output") or ""))
        model = infer_model(task, output)
        lines.extend(
            [
                f"### Agnet 调用 {index}: `{task['kind']}`",
                "",
                f"- Agnet：`{task['claimed_by']}`",
                f"- 模型：`{model}`",
                f"- 状态：`{task['status']}`",
                f"- 分数：`{task['score']}`",
                f"- 观测信号：`{observation['signal']}`",
                "",
                "#### 给模型的 user prompt 结构",
                "",
                "```text",
                user_prompt_shape_for_goal(goal),
                "```",
                "",
                "#### 接手 / 交接机制",
                "",
                handoff_description(goal, run_id, task, output),
                "",
                "#### 本次任务输入 task.input",
                "",
                "```text",
                redact_sensitive_text(str(task["input"])).strip(),
                "```",
                "",
                "#### Agnet / 模型实际输出 task.output",
                "",
                "```text",
                output.strip(),
                "```",
                "",
            ]
        )
    return lines


def redact_sensitive_text(text: str) -> str:
    replacements = [
        (r"sk-[A-Za-z0-9]{20,}", "sk-<redacted>"),
        (r"AccountKey=[^;\s`]+", "AccountKey=<redacted>"),
        (r"password=[^,;\s`]+", "password=<redacted>"),
        (r"BEGIN [A-Z ]*PRIVATE KEY", "BEGIN <redacted> PRIVATE KEY"),
    ]
    redacted = text
    for pattern, replacement in replacements:
        redacted = re.sub(pattern, replacement, redacted, flags=re.IGNORECASE)
    return redacted


def handoff_description(goal: str, run_id: str, task: dict[str, object], output: str) -> str:
    if goal.startswith("连续性长推理场景"):
        match = re.search(r"chain_edge=([^;\\n]+)", output)
        edge = match.group(1).strip() if match else "<missing>"
        marker = edge.split("->")[-1] if "->" in edge else str(task.get("kind") or "<unknown>")
        return "\n".join(
            [
                f"- 当前输出前缀记录 `chain_edge={edge}`，证明本 Agnet 承接了上一阶段。",
                "- 调用前，wrapper 会把 `Previous marker` 和 `Previous summary` 放入 user prompt。",
                f"- 执行后，wrapper 把输出摘要写入 `chain:{run_id}:{marker}:summary`。",
                f"- 同时推进 `chain:{run_id}:cursor`，并写入 `chain:{run_id}:edge:{edge}=done`。",
                "- 下一个 Agnet 读取这个 summary 和 edge 后继续执行，所以接手不是靠口头描述，而是靠共享状态字段完成。",
            ]
        )
    if goal.startswith("真实全面场景"):
        return "\n".join(
            [
                "- 每个 Agnet 从 PostgreSQL 任务池 claim 自己的子任务。",
                "- 共享状态键列表作为上下文输入，让后续任务能看到已有 run/task/agent 状态。",
                "- 完成后写回 task output、score、observation 和 Redis Stream 事件，供收敛阶段读取。",
            ]
        )
    return "\n".join(
        [
            "- Agnet 从共享任务池 claim 与自身 capability 匹配的任务。",
            "- 完成后写回 output、score、observation、heartbeat 和 pheromone 分数。",
            "- 后续 Agnet 通过 shared state、task status 和 pheromone score 感知前序执行结果。",
        ]
    )


def system_prompt_for_goal(goal: str) -> str:
    if goal.startswith("full live test"):
        return (
            "You are a minimal Agnet worker inside a swarm. "
            "Return a concise result that can be scored and converged."
        )
    if goal.startswith("真实全面场景"):
        return (
            "You are a senior coding/algorithm agent in a multi-task swarm. "
            "Return a concrete engineering answer for the assigned subtask. "
            "Do not include secrets."
        )
    if goal.startswith("连续性长推理场景"):
        return (
            "You are one stage in a continuous long-reasoning swarm. "
            "Carry forward prior conclusions, expose risks, and hand off a concise next-state. "
            "Do not reveal secrets."
        )
    return "<unknown system prompt>"


def user_prompt_shape_for_goal(goal: str) -> str:
    if goal.startswith("full live test"):
        return "\n".join(
            [
                "Task kind: <task.kind>",
                "Task input: <task.input>",
                "Known shared state keys: <sorted(shared_state.keys())>",
            ]
        )
    if goal.startswith("真实全面场景"):
        return "\n".join(
            [
                "Task kind: <task.kind>",
                "Task input:",
                "<task.input>",
                "",
                "Shared state keys: <sorted(shared_state.keys())>",
            ]
        )
    if goal.startswith("连续性长推理场景"):
        return "\n".join(
            [
                "Previous marker: <previous_marker>",
                "Previous summary:",
                "<previous_step_summary>",
                "",
                "Task kind: <task.kind>",
                "Task input:",
                "<task.input>",
            ]
        )
    return "<unknown user prompt shape>"


def infer_model(task: dict[str, object], output: str) -> str:
    for pattern in [r"used_model=([^;\\n]+)", r"primary_model=([^;\\n]+)", r"model=([^;\\n]+)"]:
        match = re.search(pattern, output)
        if match:
            return match.group(1).strip()
    match = re.search(r"model=([^;\\n]+)", str(task.get("input") or ""))
    if match:
        return match.group(1).strip()
    return "<unknown>"


def fetch_convergence(store: PostgresRedisBlobSwarmStore, run_id: str) -> dict[str, object] | None:
    def operation(cur):
        cur.execute(
            """
            select run_id, goal, accepted_score, completed_tasks, observations, artifact_path, created_at
            from swarm_convergence
            where run_id = %s
            """,
            (run_id,),
        )
        row = cur.fetchone()
        if row is None:
            return None
        return {
            "run_id": row[0],
            "goal": row[1],
            "accepted_score": row[2],
            "completed_tasks": row[3],
            "observations": row[4],
            "artifact_path": row[5],
            "created_at": row[6],
        }

    return store._run_pg(operation)


def fetch_tasks(store: PostgresRedisBlobSwarmStore, task_ids: list[str]) -> dict[str, dict[str, object]]:
    def operation(cur):
        cur.execute(
            """
            select id, kind, input, status, claimed_by, output, score
            from swarm_tasks
            where id = any(%s)
            """,
            (task_ids,),
        )
        return {
            row[0]: {
                "id": row[0],
                "kind": row[1],
                "input": row[2],
                "status": row[3],
                "claimed_by": row[4],
                "output": row[5],
                "score": row[6],
            }
            for row in cur.fetchall()
        }

    return store._run_pg(operation)


if __name__ == "__main__":
    main()