"""coding_agents_toolkit.py: a small, stdlib-only lab for Week 10 (Coding & Software Agents).

It builds on swe_bench_toolkit.py from Week 9a (same folder) and adds:
  1. Ledger and Workspace: an agent-computer interface (windowed viewer, capped search, linter-gated edit,
     a deny rule) with turn and token accounting.
  2. AciAgent: a scripted agent loop that works only through that interface.
  3. AgentlessPipeline: localize, repair (sampled), validate with a reproduction test, select by voting.
  4. Reproduction scripts and tests built from the issue text.
  5. Repo map: an ast symbol index with a link-based ranking and a token budget.
  6. Long-running harness: feature list, progress notes, checkpoints, and context-window sessions.
  7. Small helpers: the 4-tuple view of a task, an AGENTS.md reader.

Everything is deterministic. The scripted agents are programs, not language models. Every number, price,
probability and budget below is made up by me for teaching. The runner is not a security sandbox.
Run `python3 coding_agents_toolkit.py` for the demo, `python3 -m unittest test_coding_agents_toolkit` for tests.
"""
import ast
import builtins
import os
import re
import sys
from collections import Counter

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import swe_bench_toolkit as S

mean, rng_for, est_tokens = S.mean, S.rng_for, S.est_tokens
SCALE = S.CONTEXT_SCALE


# ---------------------------------------------------------------- Part 1: ledger and interface
class Ledger:
    """Counts turns and tokens. In a chat-style loop every turn re-reads the whole history, so input tokens
    grow with the square of the number of turns. `call` is for independent single-shot calls (no history)."""

    def __init__(self, tier="mid", system_tokens=600):
        self.tier, self.context, self.turns, self.tin, self.tout = tier, system_tokens, 0, 0, 0

    def turn(self, out_tokens, obs_tokens=0):
        self.tin += self.context
        self.tout += out_tokens
        self.context += out_tokens + obs_tokens
        self.turns += 1

    def call(self, prompt_tokens, out_tokens):
        self.tin += prompt_tokens
        self.tout += out_tokens
        self.turns += 1

    def cost(self, scale=SCALE, cached_frac=0.0, cache_price=0.1):
        """Dollars at the made-up tier price. cached_frac of input tokens are billed at cache_price times the rate."""
        pin, pout = S.TIERS[self.tier]
        eff_in = self.tin * ((1 - cached_frac) + cached_frac * cache_price)
        return (eff_in * pin + self.tout * pout) * scale / 1e6

    def result(self, patch, log):
        res = S.AgentResult(patch, self.tin * SCALE, self.tout * SCALE, self.turns, log)
        res.peak_context = self.context * SCALE  # how full the context window got (tokens, scaled)
        return res


LINT_MODES = ("none", "syntax", "names")


def undefined_names(tree):
    """Names that are read but never bound anywhere in the module (a coarse check: one global scope)."""
    defined = set(dir(builtins))
    for n in ast.walk(tree):
        if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            defined.add(n.name)
        elif isinstance(n, ast.arg):
            defined.add(n.arg)
        elif isinstance(n, ast.Name) and isinstance(n.ctx, (ast.Store, ast.Del)):
            defined.add(n.id)
        elif isinstance(n, (ast.Import, ast.ImportFrom)):
            for a in n.names:
                defined.add((a.asname or a.name).split(".")[0])
        elif isinstance(n, ast.ExceptHandler) and n.name:
            defined.add(n.name)
    bad = [(n.id, n.lineno) for n in ast.walk(tree)
           if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load) and n.id not in defined]
    return sorted(bad, key=lambda x: (x[1], x[0]))


def lint_source(src, mode):
    """Return an error message, or None when the source passes. mode: none, syntax (ast.parse), names."""
    if mode not in LINT_MODES:
        raise ValueError("unknown lint mode %r" % mode)
    if mode == "none":
        return None
    try:
        tree = ast.parse(src)
    except SyntaxError as e:
        return "SyntaxError: %s (line %s)" % (e.msg, e.lineno)
    if mode == "names":
        bad = undefined_names(tree)
        if bad:
            return "undefined name %s (line %d)" % bad[0]
    return None


def _lines(src):
    parts = src.split("\n")
    return parts[:-1] if src.endswith("\n") else parts


def _unlines(lines, trailing=True):
    return "\n".join(lines) + ("\n" if trailing and lines else "")


COMMANDS = ("open", "goto", "scroll_down", "scroll_up", "find_file", "search_dir", "search_file", "edit",
            "run_tests", "python", "revert", "diff", "submit")


class Workspace:
    """The agent-computer interface. Every method returns text, which is all the agent ever sees.
    The viewer shows at most `window` lines (SWE-agent: 100), search shows at most `search_cap` hits
    (SWE-agent: 50), and `edit` refuses a change that fails the linter (lint=syntax or names)."""

    def __init__(self, files, window=100, search_cap=50, lint="syntax", deny=(), ledger=None):
        self.base, self.files = dict(files), dict(files)
        self.window, self.search_cap, self.lint, self.deny = window, search_cap, lint, tuple(deny)
        self.cur, self.top = None, 1
        self.ledger = ledger or Ledger()
        self.stats = Counter()
        self.history = []
        self.last_tests = {}

    # ---- viewer
    def _view(self):
        lines = _lines(self.files[self.cur])
        n = len(lines)
        a = self.top
        b = min(n, a + self.window - 1)
        out = ["[File: %s (%d lines total)]" % (self.cur, n), "(%d more lines above)" % (a - 1)]
        out += ["%d: %s" % (i, lines[i - 1]) for i in range(a, b + 1)]
        out.append("(%d more lines below)" % (n - b))
        return "\n".join(out)

    def _place(self, line):
        n = max(1, len(_lines(self.files[self.cur])))
        top = (line if line is not None else 1) - (self.window // 2 if line is not None else 0)
        self.top = max(1, min(top, max(1, n - self.window + 1)))

    def open(self, path, line=None):
        if path not in self.files:
            return "Error: file not found: %s" % path
        self.cur = path
        self._place(line)
        return self._view()

    def goto(self, line):
        if self.cur is None:
            return "Error: no file is open. Use open first."
        n = len(_lines(self.files[self.cur]))
        if not 1 <= line <= n:
            return "Error: line %d is outside the file (1 to %d)" % (line, n)
        self._place(line)
        return self._view()

    def scroll_down(self):
        if self.cur is None:
            return "Error: no file is open. Use open first."
        n = len(_lines(self.files[self.cur]))
        self.top = max(1, min(self.top + self.window, max(1, n - self.window + 1)))
        return self._view()

    def scroll_up(self):
        if self.cur is None:
            return "Error: no file is open. Use open first."
        self.top = max(1, self.top - self.window)
        return self._view()

    # ---- search
    def find_file(self, name):
        hits = [p for p in sorted(self.files) if name in p.split("/")[-1]]
        if not hits:
            return 'No files named "%s".' % name
        if len(hits) > self.search_cap:
            return "More than %d files match. Please narrow your search." % self.search_cap
        return "Found %d file(s):\n" % len(hits) + "\n".join(hits)

    def _grep(self, term, paths):
        hits = []
        for p in paths:
            for i, ln in enumerate(_lines(self.files[p]), 1):
                if term in ln:
                    hits.append("%s:%d: %s" % (p, i, ln.strip()))
        if not hits:
            return 'No matches found for "%s".' % term
        if len(hits) > self.search_cap:
            return 'More than %d matches for "%s". Please narrow your search.' % (self.search_cap, term)
        return 'Found %d match(es) for "%s":\n' % (len(hits), term) + "\n".join(hits)

    def search_dir(self, term):
        return self._grep(term, sorted(self.files))

    def search_file(self, term, path=None):
        path = path or self.cur
        if path not in self.files:
            return "Error: file not found: %s" % str(path)
        return self._grep(term, [path])

    # ---- edit
    def edit(self, start, end, text):
        if self.cur is None:
            return "Error: no file is open. Use open first."
        if any(d in self.cur for d in self.deny):
            self.stats["blocked"] += 1
            return "Edit blocked: %s matches a protected path." % self.cur
        lines = _lines(self.files[self.cur])
        if not 1 <= start <= end <= len(lines):
            return "Error: bad line range %d to %d (the file has %d lines)" % (start, end, len(lines))
        new = text[:-1].split("\n") if text.endswith("\n") else (text.split("\n") if text else [])
        new_lines = lines[:start - 1] + new + lines[end:]
        new_src = _unlines(new_lines, self.files[self.cur].endswith("\n"))
        if self.cur.endswith(".py"):
            msg = lint_source(new_src, self.lint)
            if msg:
                self.stats["rejected"] += 1
                lo, hi = max(1, start - 2), min(len(new_lines), start + len(new) + 1)
                shown = "\n".join("%d: %s" % (i, new_lines[i - 1]) for i in range(lo, hi + 1))
                return "Your edit was not applied (%s).\nThis is how it would have looked:\n%s\nFix the edit and try again." % (msg, shown)
        self.files[self.cur] = new_src
        self.stats["edits"] += 1
        self.top = max(1, start - 3)
        return "File updated.\n" + self._view()

    # ---- run and submit
    def run_tests(self):
        res = S.run_tests(self.files)
        self.last_tests = res
        bad = sorted(t for t, s in res.items() if s != "PASS")
        head = "%d passed, %d not passing" % (len(res) - len(bad), len(bad))
        return head + ("\n" + "\n".join(bad[:5]) if bad else "")

    def python(self, expr, setup=""):
        if self.cur is None or not self.cur.endswith(".py"):
            return "Error: open a python file first."
        return "=> " + S.eval_snippet(self.files, self.cur, setup, expr)

    def revert(self, path=None):
        if path is None:
            self.files = dict(self.base)
        elif path in self.base:
            self.files[path] = self.base[path]
        else:
            self.files.pop(path, None)
        return "Reverted."

    def diff(self):
        return S.make_diff(self.base, self.files)

    def submit(self):
        return self.diff()

    def act(self, cmd, *args, think=30):
        """One agent turn: run a command, charge the ledger for the command and for the text sent back."""
        if cmd not in COMMANDS:
            raise ValueError("unknown command %r" % cmd)
        obs = getattr(self, cmd)(*args)
        out = think + est_tokens(cmd + " " + " ".join(str(a) for a in args))
        self.ledger.turn(out, est_tokens(obs))
        self.history.append((cmd, args, obs))
        return obs


# ---------------------------------------------------------------- Part 2: an agent loop on the interface
def slip_edit(text, kind):
    """Plant a typical slip in an edit. 'paren' leaves an unclosed bracket (a syntax error); 'name' mistypes the
    first parameter at its last use (passes ast.parse, fails at run time)."""
    if kind == "name":
        m = re.match(r"\s*def \w+\(((?:self, )?)(\w+)", text.split("\n")[0])
        if m:
            p = m.group(2)
            body = text.split("\n", 1)[1] if "\n" in text else ""
            spots = [x for x in re.finditer(r"\b%s\b" % re.escape(p), body)]
            if spots:
                x = spots[-1]
                body = body[:x.start()] + p + "_" + body[x.end():]
                return text.split("\n")[0] + "\n" + body
    return text.rstrip("\n") + " ("


class AciAgent:
    """Locate with BM25, open the code through the viewer, try single edits through the edit command, check
    each against the issue example and the visible tests, submit the first edit that passes.
    `slip` is the made-up chance that an edit is typed with a mistake. With lint=syntax the interface rejects
    a syntax slip and the agent retypes the edit (one more turn); with lint=none the bad edit goes in."""

    def __init__(self, name, repertoire=None, lint="syntax", slip=0.0, max_turns=60, window=15, k_units=3,
                 tier="mid", search_cap=50, use_tests=True):
        self.name, self.repertoire = name, list(repertoire or S.FULL_OPS)
        self.lint, self.slip, self.max_turns, self.window = lint, slip, max_turns, window
        self.k_units, self.tier, self.search_cap, self.use_tests = k_units, tier, search_cap, use_tests

    def solve(self, view, hint=None):
        led = Ledger(self.tier)
        ws = Workspace(view.files, window=self.window, search_cap=self.search_cap, lint=self.lint,
                       deny=("tests/",), ledger=led)
        probe = S.ScriptedAgent("probe", S.STD_OPS, k_units=self.k_units)
        units = probe._units(view)
        ranked = probe._rank(view, units, hint)[:self.k_units]
        repro = S.parse_repro(view.issue)
        log = []
        for kw in S.content_tokens(view.issue):
            if ws.search_dir(kw).startswith("Found"):
                ws.act("search_dir", kw)
                break
        base_pass = set()
        if self.use_tests:
            ws.act("run_tests")
            base_pass = {t for t, s in ws.last_tests.items() if s == "PASS"}
        tried = 0
        for key in ranked:
            path, name = key
            a, b, text = units[key]
            ws.act("open", path, a)
            if repro:
                ws.act("python", repro[1], repro[0])
            for fam, new_unit in S.candidates_for_unit(text, self.repertoire, repro[2] if repro else None):
                if led.turns >= self.max_turns:
                    log.append("turn budget used up; gave up")
                    return led.result("", log)
                tried += 1
                typed = new_unit
                if self.slip and rng_for("slip", self.name, view.id, tried).random() < self.slip:
                    kind = "paren" if rng_for("kind", self.name, view.id, tried).random() < 0.5 else "name"
                    typed = slip_edit(new_unit, kind)
                    log.append("slip (%s) on try %d" % (kind, tried))
                obs = ws.act("edit", a, b, typed, think=60)
                if obs.startswith("Your edit was not applied"):
                    log.append("edit rejected by the linter on try %d" % tried)
                    obs = ws.act("edit", a, b, new_unit, think=60)
                ok = obs.startswith("File updated")
                if ok and repro:
                    ok = ws.act("python", repro[1], repro[0]) == "=> " + repro[2]
                if ok and self.use_tests:
                    ws.act("run_tests")
                    ok = all(ws.last_tests.get(t) == "PASS" for t in base_pass)
                if ok:
                    patch = ws.act("submit")
                    log.append("submit %s edit in %s after %d tries" % (fam, name, tried))
                    return led.result(patch, log)
                ws.act("revert")
        log.append("no candidate passed after %d tries; gave up" % tried)
        return led.result("", log)


# ---------------------------------------------------------------- Part 3: reproduction scripts and tests
def repro_script(repro, import_line=""):
    """Source of a standalone reproduction script in the three-label style that the Agentless paper describes:
    it prints 'Issue reproduced', 'Issue resolved' or 'Other issues'."""
    setup, expr, expected = repro
    body = ["    " + ln for ln in setup.split("\n") if ln.strip()]
    out = ([import_line] if import_line else []) + ["", "", "def main():", "    try:"]
    out += ["    " + ln for ln in body]
    out += ["        got = %s" % expr, "    except Exception:", '        print("Other issues")', "        return",
            "    if got == %s:" % expected, '        print("Issue resolved")', "    else:", '        print("Issue reproduced")',
            "", "", "main()", ""]
    return "\n".join(out)


def run_repro(files, module_path, repro):
    """What the script would print on this version of the code (computed with the in-process runner)."""
    got = S.eval_snippet(files, module_path, repro[0], repro[1])
    if got.startswith("raises"):
        return "Other issues"
    return "Issue resolved" if got == repro[2] else "Issue reproduced"


def add_repro_test(files, test_path, repro, name="test_issue_repro"):
    """Append the issue's example to a test file as a test that asserts the expected value."""
    setup, expr, expected = repro
    body = "".join("    %s\n" % ln for ln in setup.split("\n") if ln.strip())
    src = files[test_path].rstrip("\n") + "\n\n\ndef %s():\n%s    assert (%s) == (%s)\n" % (name, body, expr, expected)
    out = dict(files)
    out[test_path] = src
    return out


def _test_path(files):
    tfs = S.test_files(files)
    return tfs[0] if tfs else None


def repro_quality(tasks):
    """For tasks whose issue has a runnable example: does the generated test fail on the base code and pass
    with the gold patch? Returns {task id: 'valid' | 'no repro' | 'passes before' | 'fails after gold'}."""
    out = {}
    for t in tasks:
        repro = S.parse_repro(t.issue)
        if repro is None:
            out[t.id] = "no repro"
            continue
        before = S.run_tests(add_repro_test(t.base, t.test_path, repro))
        gold, _ = S.apply_patch(t.base, t.gold_patch)
        after = S.run_tests(add_repro_test(gold, t.test_path, repro))
        key = t.test_path + "::test_issue_repro"
        if before.get(key) == "PASS":
            out[t.id] = "passes before"
        elif after.get(key) != "PASS":
            out[t.id] = "fails after gold"
        else:
            out[t.id] = "valid"
    return out


# ---------------------------------------------------------------- Part 4: localize, repair, validate, select
def skeleton(path, src):
    """Headers of the classes and functions in one file (the Agentless 'skeleton' idea, in my own simple form)."""
    out = [path + ":"]
    for name, (a, b, text) in S.code_units(path, src).items():
        out.append("  " + text.split("\n")[0].strip())
    return "\n".join(out)


def normalized_key(src):
    """Whitespace and comment-proof identity of a python file: the dump of its syntax tree."""
    try:
        return ast.dump(ast.parse(src))
    except SyntaxError:
        return "syntax-error:" + src


class Cand:
    def __init__(self, path, name, fam, new_src, rank):
        self.path, self.name, self.fam, self.new_src, self.rank = path, name, fam, new_src, rank
        self.key = normalized_key(new_src)


class AgentlessPipeline:
    """Fixed three-phase pipeline, no agent loop: (1) localize files then functions with BM25, (2) draw n
    sampled single-edit candidates, early candidates more often (a stand-in for sampling a model), (3) validate
    each distinct candidate against the visible tests plus a reproduction test made from the issue, then pick
    one. select: none (first sample, no validation), first (first that validates), vote (most frequent among
    those that validate; ties go to the earliest). If nothing validates the pipeline submits nothing (my choice)."""

    def __init__(self, name, repertoire=None, n_files=2, n_units=4, n_samples=8, select="vote", tier="mid", skew=2.0):
        self.name, self.repertoire = name, list(repertoire or S.FULL_OPS)
        self.n_files, self.n_units, self.n_samples, self.select, self.tier, self.skew = n_files, n_units, n_samples, select, tier, skew

    def localize(self, view, hint=None):
        probe = S.ScriptedAgent("probe", S.STD_OPS)
        units = probe._units(view)
        files = {}
        for (p, n), (a, b, text) in units.items():
            files.setdefault(p, []).append(text)
        docs = {p: S.content_tokens("\n".join(ts)) + S.content_tokens(p) for p, ts in files.items()}
        top_files = [p for _, p in S.bm25_rank(S.content_tokens(view.issue), docs)][:self.n_files]
        sub = {k: v for k, v in units.items() if k[0] in top_files}
        ranked = [k for k in probe._rank(view, sub, hint)][:self.n_units]
        return top_files, ranked, units

    def pool(self, view, ranked, units):
        repro = S.parse_repro(view.issue)
        out = []
        for key in ranked:
            path, name = key
            a, b, text = units[key]
            lines = view.files[path].split("\n")
            for fam, new_unit in S.candidates_for_unit(text, self.repertoire, repro[2] if repro else None):
                new_src = "\n".join(lines[:a - 1] + new_unit.split("\n") + lines[b:])
                out.append(Cand(path, name, fam, new_src, len(out)))
        return out

    def sample(self, view, pool, seed=0):
        """n draws from the pool. Draw 0 is the greedy one (pool[0]); later draws favour early candidates."""
        if not pool:
            return []
        rng = rng_for("sample", view.id, seed)
        draws = [pool[0]]
        for _ in range(self.n_samples - 1):
            draws.append(pool[min(len(pool) - 1, int(len(pool) * rng.random() ** self.skew))])
        return draws

    def validate(self, view, draws):
        """For each distinct candidate: (passes the visible tests that passed before, passes the reproduction test)."""
        repro = S.parse_repro(view.issue)
        tp = _test_path(view.files)
        base_pass = {t for t, s in S.run_tests(view.files).items() if s == "PASS"}
        out = {}
        for c in draws:
            if c.key in out:
                continue
            files = dict(view.files)
            files[c.path] = c.new_src
            use_repro = repro is not None and tp is not None
            if use_repro:
                files = add_repro_test(files, tp, repro)
            res = S.run_tests(files)
            reg_ok = all(res.get(t) == "PASS" for t in base_pass)
            rep_ok = res.get(tp + "::test_issue_repro") == "PASS" if use_repro else True
            out[c.key] = (reg_ok, rep_ok)
        return out

    def choose(self, draws, valid, how=None):
        how = how or self.select
        if not draws:
            return None
        if how == "none":
            return draws[0]
        good = [c for c in draws if all(valid[c.key])]
        if not good:
            return None
        if how == "first":
            return good[0]
        counts = Counter(c.key for c in good)
        best = max(counts.values())
        return next(c for c in good if counts[c.key] == best)

    def solve(self, view, hint=None, seed=0):
        led = Ledger(self.tier)
        files, ranked, units = self.localize(view, hint)
        tree = "\n".join(sorted(view.files))
        led.call(est_tokens(view.issue) + est_tokens(tree), 30)
        led.call(est_tokens(view.issue) + sum(est_tokens(skeleton(p, view.files[p])) for p in files), 40)
        led.call(est_tokens(view.issue) + sum(est_tokens(units[k][2]) for k in ranked), 40)
        pool = self.pool(view, ranked, units)
        draws = self.sample(view, pool, seed)
        ctx = est_tokens(view.issue) + sum(est_tokens(units[k][2]) for k in ranked)
        led.call(ctx, sum(est_tokens(c.new_src[:200]) for c in draws))
        repro = S.parse_repro(view.issue)
        if repro and self.select != "none":
            led.call(est_tokens(view.issue), est_tokens(repro_script(repro)))
        valid = self.validate(view, draws) if self.select != "none" else {c.key: (True, True) for c in draws}
        pick = self.choose(draws, valid)
        log = ["%d files, %d functions, pool %d, %d samples, %d distinct" % (
            len(files), len(ranked), len(pool), len(draws), len({c.key for c in draws}))]
        if pick is None:
            log.append("nothing validated; submitted nothing")
            return led.result("", log)
        patch = S.make_diff({pick.path: view.files[pick.path]}, {pick.path: pick.new_src})
        log.append("submit %s edit in %s" % (pick.fam, pick.name))
        return led.result(patch, log)


def selection_study(tasks, n_samples=8, seeds=3, skew=2.0, hidden_cache=None):
    """Draw n_samples candidates per task and seed. Pick one with each strategy. Grade the pick with the hidden
    tests. Returns {strategy: [per-task resolved share over seeds]} in task order."""
    strategies = ["greedy", "random", "visible-tests first", "visible-tests vote", "repro+tests first", "repro+tests vote", "oracle"]
    out = {s: [] for s in strategies}
    cache = hidden_cache if hidden_cache is not None else {}

    def grade(t, c):
        if c is None:
            return 0
        k = (t.id, c.key)
        if k not in cache:
            patch = S.make_diff({c.path: t.base[c.path]}, {c.path: c.new_src})
            cache[k] = int(S.evaluate(t, patch).resolved)
        return cache[k]

    for t in tasks:
        view = t.view()
        pipe = AgentlessPipeline("sel", n_samples=n_samples, skew=skew)
        top, ranked, units = pipe.localize(view, t.fn)
        pool = pipe.pool(view, ranked, units)
        acc = {s: [] for s in strategies}
        for seed in range(seeds):
            draws = pipe.sample(view, pool, seed)
            valid = pipe.validate(view, draws)
            vis = {k: (v[0], True) for k, v in valid.items()}
            rng = rng_for("rand-pick", t.id, seed)
            acc["greedy"].append(grade(t, draws[0] if draws else None))
            acc["random"].append(grade(t, rng.choice(draws) if draws else None))
            acc["visible-tests first"].append(grade(t, pipe.choose(draws, vis, "first")))
            acc["visible-tests vote"].append(grade(t, pipe.choose(draws, vis, "vote")))
            acc["repro+tests first"].append(grade(t, pipe.choose(draws, valid, "first")))
            acc["repro+tests vote"].append(grade(t, pipe.choose(draws, valid, "vote")))
            acc["oracle"].append(int(any(grade(t, c) for c in draws)))
        for s in strategies:
            out[s].append(mean(acc[s]))
    return out


# ---------------------------------------------------------------- Part 5: repo map
def symbol_index(files):
    """{path: [(kind, name, signature, line, parent)]} for classes, functions and methods of every parseable .py file."""
    out = {}
    for p in sorted(files):
        if not p.endswith(".py"):
            continue
        try:
            tree = ast.parse(files[p])
        except SyntaxError:
            continue
        syms = []
        for n in tree.body:
            if isinstance(n, ast.FunctionDef):
                syms.append(("def", n.name, "def %s(%s)" % (n.name, ast.unparse(n.args)), n.lineno, None))
            elif isinstance(n, ast.ClassDef):
                syms.append(("class", n.name, "class %s" % n.name, n.lineno, None))
                for m in n.body:
                    if isinstance(m, ast.FunctionDef):
                        syms.append(("def", m.name, "def %s(%s)" % (m.name, ast.unparse(m.args)), m.lineno, n.name))
        out[p] = syms
    return out


def reference_counts(files):
    """{path: Counter of identifiers used}: every name, attribute and imported name in the file."""
    out = {}
    for p in sorted(files):
        if not p.endswith(".py"):
            continue
        try:
            tree = ast.parse(files[p])
        except SyntaxError:
            continue
        c = Counter()
        for n in ast.walk(tree):
            if isinstance(n, ast.Name):
                c[n.id] += 1
            elif isinstance(n, ast.Attribute):
                c[n.attr] += 1
            elif isinstance(n, ast.ImportFrom):
                for a in n.names:
                    c[a.name] += 1
        out[p] = c
    return out


def file_graph(files):
    """Edges {(a, b): weight}: file a uses symbols that file b defines (and a does not define itself)."""
    idx, refs = symbol_index(files), reference_counts(files)
    defined = {p: {s[1] for s in syms} for p, syms in idx.items()}
    edges = {}
    for a in sorted(refs):
        for b in sorted(idx):
            if a == b:
                continue
            w = sum(refs[a][name] for name in defined[b] if name not in defined[a])
            if w:
                edges[(a, b)] = w
    return edges


def pagerank(nodes, edges, personal=None, d=0.85, iters=60):
    """Plain power iteration. Rank flows along edges a -> b in proportion to weight. `personal` biases the restart."""
    nodes = list(nodes)
    n = len(nodes)
    if n == 0:
        return {}
    pers = {v: (personal or {}).get(v, 0.0) for v in nodes}
    tot = sum(pers.values())
    pers = {v: (pers[v] / tot if tot else 1.0 / n) for v in nodes}
    out_w = Counter()
    for (a, b), w in edges.items():
        out_w[a] += w
    r = {v: 1.0 / n for v in nodes}
    for _ in range(iters):
        dangling = sum(r[v] for v in nodes if out_w[v] == 0)
        new = {v: (1 - d) * pers[v] + d * dangling * pers[v] for v in nodes}
        for (a, b), w in sorted(edges.items()):
            new[b] += d * r[a] * w / out_w[a]
        r = new
    return r


def repo_map(files, budget=300, focus=(), boost=5.0):
    """A compact map of the repository: files in rank order, each with its most referenced symbols, cut to a token
    budget (est_tokens). Files in `focus` get a larger restart weight, the way chat files do in Aider (the
    ranking here is my own simplified version, not Aider's). Returns (text, info)."""
    idx = symbol_index(files)
    edges = file_graph(files)
    rank = pagerank(sorted(idx), edges, personal={p: (boost if p in focus else 1.0) for p in idx})
    refs = reference_counts(files)
    uses = {}
    for p, syms in idx.items():
        for s in syms:
            uses[(p, s[1])] = sum(refs[q][s[1]] for q in refs if q != p)
    order = sorted(idx, key=lambda p: (-rank[p], p))
    chosen = {}

    def render():
        lines = []
        for p in order:
            if p not in chosen:
                continue
            lines.append(p + ":")
            for s in sorted(chosen[p], key=lambda s: s[3]):
                lines.append(("    " if s[4] else "  ") + s[2])
        return "\n".join(lines)

    for p in order:
        cand = sorted(idx[p], key=lambda s: (-uses[(p, s[1])], s[3]))
        for s in cand:
            chosen.setdefault(p, [])
            if s[4] and not any(x[1] == s[4] and x[0] == "class" for x in chosen[p]):
                head = next((x for x in idx[p] if x[0] == "class" and x[1] == s[4]), None)
                trial = chosen[p] + ([head] if head else []) + [s]
            else:
                trial = chosen[p] + [s]
            saved = chosen[p]
            chosen[p] = trial
            if est_tokens(render()) > budget:
                chosen[p] = saved
                break
        if not chosen.get(p):
            chosen.pop(p, None)
    text = render()
    total = sum(len(v) for v in idx.values())
    shown = sum(len(v) for v in chosen.values())
    return text, {"files": len(chosen), "files_total": len(idx), "symbols": shown, "symbols_total": total,
                  "tokens": est_tokens(text), "full_tokens": sum(est_tokens(files[p]) for p in idx), "rank": rank}


def synthetic_repo(n_modules=12, per_module=3, seed=0):
    """A made-up package where early modules are used by many later ones (hubs). Every function is valid python."""
    rng = rng_for("synthetic", n_modules, per_module, seed)
    files = {"pkg/__init__.py": ""}
    for i in range(n_modules):
        imports, body = [], []
        for j in range(per_module):
            calls = []
            if i:
                for _ in range(rng.choice([1, 2])):
                    k = min(i - 1, int(i * rng.random() ** 3))
                    f = "f%d_%d" % (k, rng.randrange(per_module))
                    line = "from pkg.m%02d import %s" % (k, f)
                    if line not in imports:
                        imports.append(line)
                    calls.append(f)
            expr = " + ".join("%s(x)" % c for c in calls) if calls else "x"
            body.append("def f%d_%d(x):\n    return %s + %d\n" % (i, j, expr, j))
        files["pkg/m%02d.py" % i] = "\n".join(imports) + ("\n\n\n" if imports else "") + "\n\n".join(body)
    return files


def localize_study(tasks, mask_names=False):
    """File-level top-1 and function-level top-3 hit rates, comparing an index of full function bodies with an index
    of signatures only (what a tight map would show). One merged repo holds all five toy packages. With
    mask_names the issue text has the buggy function's name replaced by the word 'it' (a harder, behaviour-only issue)."""
    merged = {}
    for t in tasks:
        for p, src in t.base.items():
            merged.setdefault(p, src)
    acc = Counter()
    for t in tasks:
        issue = re.sub(r"\b%s\b" % re.escape(t.fn), "it", t.issue) if mask_names else t.issue
        view = S.TaskView(t.id, issue, merged)
        probe = S.ScriptedAgent("probe", S.STD_OPS)
        units = probe._units(view)
        q = S.content_tokens(issue)
        by_file = {}
        for (p, n), (a, b, text) in units.items():
            by_file.setdefault(p, []).append((n, text))
        for mode in ("bodies", "signatures"):
            def doc(text):
                return S.content_tokens(text if mode == "bodies" else text.split("\n")[0])
            fdocs = {p: [w for n, tx in v for w in doc(tx)] for p, v in by_file.items()}
            top_file = S.bm25_rank(q, fdocs)[0][1]
            acc[mode + " file@1"] += top_file == t.module
            udocs = {k: doc(v[2]) + 3 * [x for x in S.tokens(k[1]) if x not in S.STOP] for k, v in units.items()}
            top3 = [k for _, k in S.bm25_rank(q, udocs)][:3]
            acc[mode + " function@3"] += any(k[1].split(".")[-1] == t.fn for k in top3)
    return {k: v / len(tasks) for k, v in acc.items()}


# ---------------------------------------------------------------- Part 6: long-running harness
def project_files(repo):
    """One package with every planted bug of that repo at once, plus the basic tests (which still pass) and one
    regression test per feature (which fail for the bugs). Returns (files, features)."""
    r = S.REPOS[repo]
    specs = [s for s in S.SPECS if s["repo"] == repo]
    src = r["src"]
    for s in specs:
        try:
            src = S._subs(src, s["bug"])
        except ValueError:
            pass
    files = S.repo_files(repo, [s["fn"] for s in specs])
    files[r["module"]] = src
    tests = files[r["test"]] + "".join("\n\ndef test_%s_regress():\n    %s\n" % (s["fn"], S.REGRESS[s["fn"]]) for s in specs)
    files[r["test"]] = tests
    res = S.run_tests(files)
    feats = []
    for s in specs:
        key = r["test"] + "::test_%s_regress" % s["fn"]
        feats.append({"id": s["fn"], "issue": s["issue"], "test": key, "passes": res.get(key) == "PASS"})
    return files, feats


class HarnessConfig:
    """Which pieces of the harness are switched on, and the made-up limits of the toy world."""

    def __init__(self, feature_list=True, progress=True, commits=True, session_context=400000, crash=0.3,
                 max_sessions=8, max_attempts=2):
        self.feature_list, self.progress, self.commits = feature_list, progress, commits
        self.session_context, self.crash, self.max_sessions, self.max_attempts = session_context, crash, max_sessions, max_attempts


def _visible_pass(files, feats):
    res = S.run_tests(files)
    hidden = {f["test"] for f in feats}
    return all(s == "PASS" for t, s in res.items() if t not in hidden) and bool(res)


def run_long(repo, cfg, seed=0, agent=None):
    """Work through a project in several sessions. A session ends when its context window (session_context,
    made-up tokens) is used up, or when the toy 'crash' cuts it off in the middle of an edit, which leaves a
    half-written file in the working tree. What survives between sessions depends on cfg."""
    agent = agent or AciAgent("harness-agent", tier="mid")
    files, feats = project_files(repo)
    pin, pout = S.TIERS[agent.tier]
    commits, progress, attempts = [dict(files)], [], Counter()
    tokens_in = tokens_out = 0
    log, sessions, declared_done = [], 0, False
    for sess in range(cfg.max_sessions):
        sessions += 1
        if cfg.progress:
            notes = "\n".join(progress) + "\n" + "\n".join("%s %s" % (f["id"], f["passes"]) for f in feats)
            orient = est_tokens(notes) * SCALE
        else:
            orient = sum(est_tokens(v) for v in files.values()) * SCALE
        used = orient
        tokens_in += orient
        if cfg.commits:
            files = dict(commits[-1])
        if cfg.feature_list:
            todo = [f for f in feats if not f["passes"] and (not cfg.progress or attempts[f["id"]] < cfg.max_attempts)]
        else:
            if _visible_pass(files, feats):
                declared_done = True
                log.append("session %d: visible tests pass, declared done" % sess)
                break
            log.append("session %d: visible tests fail and nothing says why; stopped" % sess)
            break
        if not todo:
            log.append("session %d: nothing left to try" % sess)
            break
        for f in todo:
            if used >= cfg.session_context:
                break
            res = agent.solve(S.TaskView(f["id"], f["issue"], files))
            attempts[f["id"]] += 1
            used += res.peak_context
            tokens_in += res.tokens_in
            tokens_out += res.tokens_out
            if rng_for("crash", repo, seed, sess, f["id"], attempts[f["id"]]).random() < cfg.crash:
                if res.patch:
                    mod = S.REPOS[repo]["module"]
                    try:
                        new, _ = S.apply_patch(files, res.patch)
                        files = dict(files)
                        files[mod] = new[mod][:len(new[mod]) // 2]
                    except S.PatchError:
                        pass
                log.append("session %d: cut off while working on %s" % (sess, f["id"]))
                break
            if res.patch:
                try:
                    files, _ = S.apply_patch(files, res.patch)
                except S.PatchError:
                    pass
            if S.run_tests(files).get(f["test"]) == "PASS":
                f["passes"] = True
                if cfg.commits:
                    commits.append(dict(files))
                if cfg.progress:
                    progress.append("done " + f["id"])
            else:
                if cfg.commits:
                    files = dict(commits[-1])
                if cfg.progress:
                    progress.append("failed " + f["id"])
    final = S.run_tests(files)
    ok = sum(1 for f in feats if final.get(f["test"]) == "PASS")
    start_pass = sum(1 for f in project_files(repo)[1] if f["passes"])
    return {"repo": repo, "features": len(feats), "passing": ok, "started_passing": start_pass, "sessions": sessions,
            "tokens": tokens_in + tokens_out, "cost": (tokens_in * pin + tokens_out * pout) / 1e6,
            "healthy": all(s == "PASS" for t, s in final.items() if t not in {f["test"] for f in feats}),
            "declared_done": declared_done, "log": log}


# ---------------------------------------------------------------- Part 7: small helpers
def four_tuple(task):
    """The four parts of a benchmark item (request, environment, stopping criteria, scorer) for one task."""
    return {"request": task.issue.split("\n")[0],
            "environment": "repo at the base commit, %d files" % len(task.base),
            "stopping": "the agent submits a patch, or runs out of turns or budget",
            "scorer": "%d FAIL_TO_PASS and %d PASS_TO_PASS tests must all pass" % (len(task.f2p), len(task.p2p))}


def parse_agents_md(text):
    """Read an AGENTS.md style file: '## Heading' sections with '- bullet' lines. Returns {heading: [bullets]}, plus
    the backticked commands found anywhere under the key '_commands'."""
    out, cur = {}, None
    for ln in text.split("\n"):
        if ln.startswith("#"):
            cur = ln.lstrip("#").strip()
            out.setdefault(cur, [])
        elif ln.strip().startswith("- ") and cur is not None:
            out[cur].append(ln.strip()[2:].strip())
    out["_commands"] = re.findall(r"`([^`]+)`", text)
    return out


# ---------------------------------------------------------------- Part 8: demo
def _run_agent(agent, tasks):
    out = []
    for t in tasks:
        res = agent.solve(t.view(), hint=t.fn)
        out.append((S.evaluate(t, res.patch).resolved, res.turns, S.cost_of(agent, res)))
    return out


def _row(label, runs):
    k, n = sum(r[0] for r in runs), len(runs)
    lo, hi = S.wilson(k, n)
    cost = mean(r[2] for r in runs)
    cps = "%.3f" % (sum(r[2] for r in runs) / k) if k else "n/a"
    return "%-26s %2d/%d (95%% CI %2.0f to %2.0f%%) %6.1f %9.3f %9s" % (label, k, n, 100 * lo, 100 * hi, mean(r[1] for r in runs), cost, cps)


def demo():
    tasks = S.curated_tasks()
    syn, _ = S.synthesize_tasks(limit=40)
    pool = tasks + syn
    hdr = "%-26s %-26s %6s %9s %9s" % ("system", "resolved (Wilson CI)", "turns", "$/task", "$/solved")

    print("== 1. One task through the interface (%s)" % tasks[2].id)
    t = tasks[2]
    ws = Workspace(t.base, window=8, search_cap=3, lint="syntax", deny=("tests/",))
    a, b, _ = S.code_units(t.module, t.base[t.module])["startswith_any"]
    fix = "    return any(s.startswith(p) for p in prefixes)"
    for cmd, args in [("search_dir", ("prefixes",)), ("open", (t.module, a)), ("python", ("startswith_any('foobar', ['foo'])",)),
                      ("edit", (a + 1, b, "    return (")), ("edit", (a + 1, b, fix)),
                      ("python", ("startswith_any('foobar', ['foo'])",)), ("search_dir", ("def",))]:
        obs = ws.act(cmd, *args)
        print("> %s %s\n%s" % (cmd, ", ".join(repr(x) for x in args), "\n".join("    " + x for x in obs.split("\n")[:11])))
    ws.open(t.test_path)
    print("> edit on a test file ->", ws.act("edit", 1, 1, "pass"))
    print("turns %d, input tokens %d, output tokens %d (unscaled), stats %s" % (ws.ledger.turns, ws.ledger.tin, ws.ledger.tout, dict(ws.stats)))

    print("\n== 2. The four parts of a task (request, environment, stopping criteria, scorer)")
    for k, v in four_tuple(t).items():
        print("  %-12s %s" % (k, v))

    print("\n== 3. Interface ablations on %d tasks (20 curated + 40 synthesized). Slip chance 0.4 is made up." % len(pool))
    print(hdr)
    for label, ag in [("lint syntax, no slips", AciAgent("a1")),
                      ("lint none, slips", AciAgent("a2", lint="none", slip=0.4)),
                      ("lint syntax, slips", AciAgent("a3", lint="syntax", slip=0.4)),
                      ("lint names, slips", AciAgent("a4", lint="names", slip=0.4))]:
        print(_row(label, _run_agent(ag, pool)))
    print("viewer window (lines) vs cost on the 20 curated tasks:")
    for w in (5, 15, 100):
        rs = _run_agent(AciAgent("w", window=w), tasks)
        print("  window %3d: resolved %d/20, $/task %.3f" % (w, sum(r[0] for r in rs), mean(r[2] for r in rs)))

    print("\n== 4. Agent loop vs fixed pipeline (same %d tasks)" % len(pool))
    print(hdr + "   (turns = model calls for the pipeline)")
    print(_row("agent loop (ACI)", _run_agent(AciAgent("loop"), pool)))
    for sel in ("none", "first", "vote"):
        print(_row("pipeline, select=" + sel, _run_agent(AgentlessPipeline("p-" + sel, select=sel), pool)))
    print("more samples (select=vote):")
    for n in (1, 4, 8, 16, 32):
        rs = _run_agent(AgentlessPipeline("n%d" % n, n_samples=n), pool)
        print("  n=%2d: resolved %d/%d, $/task %.3f" % (n, sum(r[0] for r in rs), len(rs), mean(r[2] for r in rs)))

    print("\n== 5. Reproduction scripts")
    rq = repro_quality(tasks)
    print("generated test fails before and passes after the gold patch: %s" % dict(Counter(rq.values())))
    repro = S.parse_repro(t.issue)
    print(repro_script(repro, "from textkit.text import *"))
    print("script on the buggy code: %s; on the gold patch: %s" % (
        run_repro(t.base, t.module, repro), run_repro(S.apply_patch(t.base, t.gold_patch)[0], t.module, repro)))

    print("== 6. Choosing among sampled patches (curated 20, 5 seeds)")
    for skew, n in ((2.0, 8), (1.0, 16)):
        st = selection_study(tasks, n_samples=n, seeds=5, skew=skew)
        print("n=%d samples, skew %.1f (1.0 is uniform; larger favours early candidates)" % (n, skew))
        for k, v in st.items():
            lo, hi = S.bootstrap_ci(v)
            print("  %-20s %5.1f of 20 resolved (task bootstrap %2.0f to %2.0f%%)" % (k, sum(v), 100 * lo, 100 * hi))

    print("\n== 7. Repo map")
    syn_repo = synthetic_repo()
    for b in (40, 120, 400):
        text, info = repo_map(syn_repo, budget=b)
        print("  budget %3d: %2d/%2d files, %2d/%2d symbols, %3d of %3d tokens" % (
            b, info["files"], info["files_total"], info["symbols"], info["symbols_total"], info["tokens"], info["full_tokens"]))
    _, info = repo_map(syn_repo, budget=400)
    top = sorted(info["rank"].items(), key=lambda x: -x[1])[:3]
    print("  most central files:", ", ".join("%s %.2f" % kv for kv in top))
    merged = {}
    for x in tasks:
        merged.update(x.base)
    text, _ = repo_map(merged, budget=120, focus=(t.module,))
    print("  map with focus on %s (budget 120):\n%s" % (t.module, "\n".join("    " + x for x in text.split("\n"))))
    for masked in (False, True):
        print("  localization, issue names the function = %s: %s" % (not masked, {k: round(v, 2) for k, v in localize_study(tasks, masked).items()}))

    print("\n== 8. Long-running harness (5 packages x 5 seeds, crash chance 0.3 per attempt, all made up)")
    cfgs = [("bare (no list, notes or commits)", dict(feature_list=False, progress=False, commits=False)),
            ("feature list only", dict(feature_list=True, progress=False, commits=False)),
            ("list + progress notes", dict(feature_list=True, progress=True, commits=False)),
            ("list + notes + commits", dict(feature_list=True, progress=True, commits=True))]
    print("%-34s %8s %9s %9s %8s" % ("setup", "passing", "sessions", "$/run", "healthy"))
    for label, kw in cfgs:
        rows = [run_long(r, HarnessConfig(**kw), seed=s) for r in S.REPOS for s in range(5)]
        print("%-34s %3d/%3d %9.1f %9.2f %5d/25" % (label, sum(r["passing"] for r in rows), sum(r["features"] for r in rows),
                                                    mean(r["sessions"] for r in rows), mean(r["cost"] for r in rows),
                                                    sum(r["healthy"] for r in rows)))
    print("full setup as the crash chance grows:")
    for c in (0.0, 0.3, 0.6):
        rows = [run_long(r, HarnessConfig(crash=c), seed=s) for r in S.REPOS for s in range(5)]
        print("  crash %.1f: passing %d/%d, sessions %.1f" % (c, sum(r["passing"] for r in rows), sum(r["features"] for r in rows),
                                                          mean(r["sessions"] for r in rows)))

    print("\n== 9. Cost and turns")
    led = Ledger("mid")
    cum = []
    for i in range(1, 31):
        led.turn(60, 250)
        cum.append(led.tin)
    print("  input tokens after 10, 20, 30 turns: %d, %d, %d (unscaled); the 30 turn run costs %.1f times the 10 turn run" % (
        cum[9], cum[19], cum[29], cum[29] / cum[9]))
    for cf in (0.0, 0.5, 0.9):
        print("  30 turns, %2.0f%% of input billed as cached at 0.1x: $%.2f" % (100 * cf, led.cost(cached_frac=cf)))

    print("\n== 10. An AGENTS.md file")
    md = "# AGENTS.md\n## Setup\n- Install: `pip install -e .`\n- Test: `python -m unittest`\n## Style\n- Keep functions short\n"
    p = parse_agents_md(md)
    print("  sections: %s; commands: %s" % ([k for k in p if k != "_commands"], p["_commands"]))


if __name__ == "__main__":
    demo()