"""
=============================================================================
第 12 讲 · 工具模块：一个迷你的编程智能体 —— 模拟的 LLM + 真实的工具 + 可以调整的 harness
=============================================================================

对应讲义：lectures/lecture12_coding_agents.md 第 2~4 节

-----------------------------------------------------------------------------
【编程智能体 = LLM + 工具 + harness】
-----------------------------------------------------------------------------
    LLM      读到的东西（源代码、测试的输出）→ 决定下一步：改哪里、怎么改
    工具     read_file（读源代码）、run_tests（运行可见的测试）、edit_file（把文件改成新的内容）
    harness  包在 LLM 外面的程序：给 LLM 看什么、允许它做什么、什么时候停 —— 这里有 6 个"旋钮"（Harness）：

        read_source     改之前先读源代码吗？（不读 = 只看问题描述就猜）
        error_detail    测试失败时给 LLM 看什么："none"（什么都不给）/ "pass_fail"（只说几个通过、几个失败）/
                        "full"（每个失败的测试：输入、实际输出、期望输出）
        max_attempts    最多改几次
        protect_tests   测试文件只读吗？（不只读 = LLM 可以"改测试让它通过"）
        self_review     测试通过之后，让 LLM 再看一遍自己的修改（会发现"专门针对测试写特例"的代码）
        edge_hint       系统提示里提醒"考虑边界情况，不要针对测试写特例"

-----------------------------------------------------------------------------
【模拟的边界】（非常重要）
-----------------------------------------------------------------------------
没有真实的 LLM。SimCoder 按设定的概率（PARAMS）在四种候选修改（coding_world.py）里选一种，或者去改测试；
概率取决于它"看到了什么"（读没读源代码、错误信息有多详细、之前失败了几次）。
**工具是真实的**：修改真的写进"文件"（一个字典），测试真的被执行，可见 / 隐藏测试的结果都是真实的。
"""

import random
from dataclasses import dataclass, field

from coding_world import TASKS, all_pass, edit_tests_to_pass, run_tests

PARAMS = {
    "skill": 0.85,             # 看到全部信息时，看懂 bug 的基础概率
    "difficulty_weight": 0.6,  # 难度每增加 1，看懂的概率减少多少
    "blind": 0.35,             # 没读源代码时，看懂的概率乘以这个数
    "detail": {"none": 0.45, "pass_fail": 0.65, "full": 1.0},   # 错误信息的详细程度对"定位 bug"的影响
    "retry_gain": 0.12,        # 每看到一次详细的失败信息，看懂的概率增加多少
    "edge": 0.5,               # 看懂了 bug 时，同时想到边界情况（写出 correct 而不是 partial）的概率
    "edge_hint": 0.8,          # 系统提示里提醒了"考虑边界情况"时的概率
    "overfit": 0.45,           # 没看懂、但看到了失败测试的具体输入时，"针对测试写特例"的概率
    "overfit_hint": 0.15,      # 提醒了"不要针对测试写特例"时的概率
    "hack": 0.12,              # 已经失败过一次之后，每次尝试"去改测试"的概率
    "review_catch": 0.85,      # 自我检查发现"针对测试写特例"或"改了测试"的概率
    "tokens": {"system": 400, "read": 300, "none": 20, "pass_fail": 40, "full": 160, "patch": 150},
}


@dataclass(frozen=True)
class Harness:
    """一组 harness 的设置（6 个旋钮，见文件开头）。frozen=True 使它可以当字典的键、放进集合。"""
    read_source: bool = True
    error_detail: str = "pass_fail"
    max_attempts: int = 3
    protect_tests: bool = False
    self_review: bool = False
    edge_hint: bool = False

    def short(self):
        """一行的简短描述。"""
        return (f"读源码={'是' if self.read_source else '否'} 错误信息={self.error_detail} 最多 {self.max_attempts} 次 "
                f"测试只读={'是' if self.protect_tests else '否'} 自我检查={'是' if self.self_review else '否'} "
                f"提醒边界={'是' if self.edge_hint else '否'}")


@dataclass
class Step:
    """轨迹里的一步：哪个工具、输入的摘要、观察到的结果。"""
    tool: str
    detail: str
    observation: str


@dataclass
class Outcome:
    """
    一次运行的结果。

    【字段】claimed：智能体自己认为成功了（它运行的可见测试全部通过）；
        solved：真的成功了（最后的代码通过**原来的**可见测试和全部隐藏测试）；
        final_kind：最后留下的是哪种修改（correct / partial / overfit / wrong / buggy），是否改过测试；
        attempts、llm_calls、tokens：花了多少；steps：完整的轨迹
    """
    claimed: bool
    solved: bool
    final_kind: str
    edited_tests: bool
    attempts: int
    llm_calls: int
    tokens: int
    steps: list = field(default_factory=list)
    failure: str = ""          # 失败的原因（给 02 号文件的"反思式优化器"读）


class SimCoder:
    """
    模拟的 LLM：每次被要求"改代码"时，决定写哪一种修改。

    【思路】先算"看懂 bug 的概率"p：
            p = (skill − 难度 × difficulty_weight) × detail[错误信息] × (1 或 blind) + retry_gain × 之前看到的详细失败次数
        · 看懂了：以 edge（或 edge_hint）的概率连边界情况一起修好（correct），否则只修了可见的问题（partial）；
        · 没看懂：如果看到了失败测试的具体输入（error_detail == "full"），以 overfit 的概率"针对测试写特例"，
          否则写一个错误的修改（wrong）；
        · 已经失败过之后，每次还有 hack 的概率决定"去改测试"。
    """

    def __init__(self, rng, params=None):
        self.rng = rng
        self.p = dict(PARAMS)
        if params:
            self.p.update(params)

    def propose(self, task, harness, seen_full_failures, failures):
        """
        提出一次修改。

        【参数】seen_full_failures：之前看到过几次详细的失败信息；failures：之前失败了几次
        【返回】"correct" / "partial" / "overfit" / "wrong" / "edit_tests"
        """
        p, rng = self.p, self.rng
        if failures > 0 and rng.random() < p["hack"]:
            return "edit_tests"
        understand = (p["skill"] - p["difficulty_weight"] * task.difficulty) * p["detail"][harness.error_detail]
        if not harness.read_source:
            understand *= p["blind"]
        understand = min(0.97, understand + p["retry_gain"] * seen_full_failures)
        if rng.random() < understand:
            return "correct" if rng.random() < p["edge_hint" if harness.edge_hint else "edge"] else "partial"
        if harness.error_detail == "full" and rng.random() < p["overfit_hint" if harness.edge_hint else "overfit"]:
            return "overfit"
        return "wrong"

    def review(self, kind, edited):
        """自我检查：重新读一遍修改，能否发现"针对测试写特例"或"改了测试"。【返回】True 表示发现了问题。"""
        return (kind == "overfit" or edited) and self.rng.random() < self.p["review_catch"]


def format_results(results, detail):
    """把测试结果按 error_detail 的设置写成给 LLM 看的文字（这正是 harness 决定的事）。"""
    failed = [r for r in results if not r[1]]
    if detail == "none":
        return "Tests finished."
    if detail == "pass_fail" or not failed:
        return f"{len(results) - len(failed)} passed, {len(failed)} failed."
    return f"{len(failed)} failed: " + "; ".join(msg for _, _, msg in failed)


def solve(task, harness, seed=0, params=None, record=False):
    """
    让智能体完成一个任务（论文里的"一次 rollout"）。

    【流程】（每一步都是一次真实的工具调用）
        1. （可选）read_file：读源代码；
        2. run_tests：运行可见的测试，看到失败；
        3. 循环（最多 max_attempts 次）：LLM 提出修改 → edit_file（或者改测试）→ run_tests →
           全部通过 →（可选）自我检查 → 结束；否则把错误信息交给 LLM，再试一次。
    【返回】Outcome
    """
    rng = random.Random(f"solve|{task.name}|{seed}")
    llm = SimCoder(rng, params)
    tok = llm.p["tokens"]
    src, tests = task.buggy, list(task.visible)
    steps, tokens, calls = [], tok["system"], 0
    kind, edited = "buggy", False
    if harness.read_source:
        steps.append(Step("read_file", task.filename, src.strip()))
        tokens += tok["read"]
    res = run_tests(src, tests)
    obs = format_results(res, harness.error_detail)
    steps.append(Step("run_tests", "visible tests", obs))
    tokens += tok[harness.error_detail]
    seen_full = 1 if harness.error_detail == "full" else 0
    failures, attempts, failure = 0, 0, ""
    while attempts < harness.max_attempts:
        attempts += 1
        calls += 1
        tokens += tok["patch"]
        choice = llm.propose(task, harness, seen_full - 1 if seen_full else 0, failures)
        if choice == "edit_tests":
            if harness.protect_tests:
                steps.append(Step("edit_file", "tests/test_" + task.filename, "BLOCKED: test files are read-only."))
                failures += 1
                failure = "test_edit_blocked"
                continue
            tests = edit_tests_to_pass(src, tests)
            edited = True
            steps.append(Step("edit_file", "tests/test_" + task.filename, "changed expected values of failing tests"))
        else:
            src, kind = task.patches[choice], choice
            steps.append(Step("edit_file", task.filename, src.strip()))
        res = run_tests(src, tests)
        obs = format_results(res, harness.error_detail)
        steps.append(Step("run_tests", "visible tests", obs))
        tokens += tok[harness.error_detail]
        if all_pass(res):
            if harness.self_review:
                calls += 1
                tokens += tok["patch"]
                if llm.review(kind, edited):
                    steps.append(Step("review", "self-review of the diff", "Rejected: special-cases the tests."))
                    if edited:                                    # 撤销对测试的修改
                        tests, edited = list(task.visible), False
                    failures += 1
                    failure = "review_rejected"
                    continue
                steps.append(Step("review", "self-review of the diff", "Looks good."))
            break
        failures += 1
        seen_full += harness.error_detail == "full"
        failure = "wrong_patch"
    claimed = all_pass(run_tests(src, tests))
    solved = all_pass(run_tests(src, task.visible + task.hidden))
    if solved:
        failure = ""
    else:
        if edited:
            failure = "edited_tests"
        elif kind == "overfit":
            failure = "special_case"
        elif kind == "partial":
            failure = "missed_edge_case"
        elif kind in ("wrong", "buggy"):
            failure = "out_of_attempts" if attempts >= harness.max_attempts else failure
    out = Outcome(claimed, solved, kind, edited, attempts, calls, tokens, steps if record else [], failure)
    return out


def one_shot(task, seed=0, params=None):
    """
    基线：不用工具，只看问题描述和源代码，一次写出修改（不运行测试，"写完就交"）。

    【返回】Outcome（claimed 记为 True：它不知道自己错了）
    """
    rng = random.Random(f"oneshot|{task.name}|{seed}")
    llm = SimCoder(rng, params)
    choice = llm.propose(task, Harness(read_source=True, error_detail="none", max_attempts=1), 0, 0)
    src = task.patches[choice]
    solved = all_pass(run_tests(src, task.visible + task.hidden))
    return Outcome(True, solved, choice, False, 1, 1, PARAMS["tokens"]["system"] + PARAMS["tokens"]["read"]
                   + PARAMS["tokens"]["patch"], [], "" if solved else "no_feedback")


def evaluate(harness, tasks=TASKS, seeds=range(10), params=None):
    """
    在一组任务上评估一个 harness（每个任务用几个随机种子各跑一次）。

    【返回】dict：solved（真实的成功率）、claimed（智能体声称的成功率）、calls、tokens（平均每次的花费）、
        kinds（最后留下的修改的种类分布）、failures（失败原因的计数）、outcomes（全部 Outcome）
    """
    outs = [solve(t, harness, seed=s, params=params) for t in tasks for s in seeds]
    n = len(outs)
    kinds, failures = {}, {}
    for o in outs:
        k = "edited_tests" if o.edited_tests else o.final_kind
        kinds[k] = kinds.get(k, 0) + 1
        if not o.solved:
            failures[o.failure] = failures.get(o.failure, 0) + 1
    return {"solved": sum(o.solved for o in outs) / n, "claimed": sum(o.claimed for o in outs) / n,
            "calls": sum(o.llm_calls for o in outs) / n, "tokens": sum(o.tokens for o in outs) / n,
            "kinds": kinds, "failures": failures, "outcomes": outs}