#!/usr/bin/env python3
"""agentic_harness.py -- PHASE 1: a deterministic long-context agentic harness (audit 251).

WHY: "did the agent do well?" is not a measurement. Every turn is machine-checked so a failure names its
own layer (tool-call degradation, context loss, wasteful rereading, loops, error-swallowing, hallucination,
planning drift) instead of merely "feeling frustrating".

The environment is synthetic and in-memory: a small repo whose tests are really executed, plus a large
distractor spec so the context grows like a real session. Tools are executed client-side.

Per turn: prompt_tokens, cached_tokens, completion_tokens, wall_s, finish_reason, loop_abort,
think_reentry, tool calls, and any assertion violations.
Output: /var/tmp/agentic/<task>.jsonl and /var/tmp/agentic/summary.json

usage: agentic_harness.py [T1 T2 T3 T4]
"""
from __future__ import annotations

import json, os, re, sys, time, urllib.request

URL = "http://127.0.0.1:8300/v1/chat/completions"
OUT = "/var/tmp/agentic"
MODEL = "glm-5.3-exl3"
TEMP = 0.6
THINK_BUDGET = 3000
MAX_TOKENS = 4000
MAX_TURNS = 22

CALC = '"""Small arithmetic helpers."""\n\n\ndef add(a, b):\n    return a - b   # BUG: should be a + b\n\n\ndef mul(a, b):\n    return a * b\n\n\ndef describe():\n    return "calc module"\n'
UTIL = "def clamp(x, lo, hi):\n    return max(lo, min(hi, x))\n\n\ndef is_even(n):\n    return n % 2 == 0\n"
T_CALC = ("def test_add():\n    assert mod.add(2, 3) == 5, \"add(2,3) should be 5\"\n\n\n"
          "def test_mul():\n    assert mod.mul(3, 4) == 12\n\n\n"
          "def test_describe():\n    assert mod.describe() == \"calc module\"\n")
T_UTIL = ("def test_clamp():\n    assert mod.clamp(5, 0, 3) == 3\n\n\n"
          "def test_is_even():\n    assert mod.is_even(4) is True\n")
SPEC = "# Pipeline specification\n\n" + "\n\n".join(
    "## Section %d\nThe pipeline validates incoming records, applies transformation stage %d, and emits "
    "structured output. Invariant %d: ordering is preserved and ids are stable across restarts. "
    "See ticket ENG-%d." % (i, i, i, 2000 + i) for i in range(140))

REPO = {"src/calc.py": CALC, "src/util.py": UTIL, "tests/test_calc.py": T_CALC,
        "tests/test_util.py": T_UTIL, "docs/SPEC.md": SPEC}
MODS = {"src/calc.py": "calc", "src/util.py": "util"}
STATES = dict(REPO)


def _tool(name, desc, props=None, req=None):
    return {"type": "function", "function": {"name": name, "description": desc,
            "parameters": {"type": "object", "properties": props or {}, "required": req or []}}}


TOOLS = [
    _tool("list_files", "List every file path in the repo."),
    _tool("read_file", "Read a file's full contents.", {"path": {"type": "string"}}, ["path"]),
    _tool("write_file", "Replace a file's full contents.",
          {"path": {"type": "string"}, "content": {"type": "string"}}, ["path", "content"]),
    _tool("run_tests", "Run the test suite; returns per-test results."),
    _tool("search", "Search file contents for a literal string.", {"query": {"type": "string"}}, ["query"]),
    _tool("fail_always", "Diagnostic tool that always returns an error."),
]


def run_tests_report() -> str:
    import types
    mod = types.ModuleType("mod")
    for path in MODS:
        try:
            exec(compile(STATES[path], path, "exec"), mod.__dict__)
        except Exception as e:
            return "IMPORT ERROR in %s: %s: %s" % (path, type(e).__name__, e)
    tns = {"mod": mod}
    for tf in ("tests/test_calc.py", "tests/test_util.py"):
        try:
            exec(compile(STATES[tf], tf, "exec"), tns)
        except Exception as e:
            return "TEST COLLECTION ERROR: %s: %s" % (type(e).__name__, e)
    lines, npass, nfail = [], 0, 0
    for name in sorted(k for k in tns if k.startswith("test_")):
        fn = tns[name]
        if not callable(fn):
            continue
        try:
            fn()
            lines.append("PASS: " + name)
            npass += 1
        except AssertionError as e:
            lines.append("FAIL: %s -- %s" % (name, e or "assertion failed"))
            nfail += 1
        except Exception as e:
            lines.append("FAIL: %s -- %s: %s" % (name, type(e).__name__, e))
            nfail += 1
    return "\n".join(lines) + "\n\n%d passed, %d failed" % (npass, nfail)


def exec_tool(name, args):
    if name == "list_files":
        return "REPO FILES:\n" + "\n".join(sorted(STATES))
    if name == "read_file":
        p = args.get("path", "")
        return STATES[p] if p in STATES else "ERROR: no such file: %s" % p
    if name == "write_file":
        p = args.get("path", "")
        if p not in STATES:
            return "ERROR: no such file: %s" % p
        STATES[p] = args.get("content", "")
        return "OK: wrote %d bytes to %s" % (len(STATES[p]), p)
    if name == "run_tests":
        return run_tests_report()
    if name == "search":
        q = args.get("query", "")
        hits = ["%s:%d: %s" % (p, i + 1, ln.strip()) for p, c in sorted(STATES.items())
                for i, ln in enumerate(c.splitlines()) if q and q in ln]
        return ("MATCHES:\n" + "\n".join(hits[:40])) if hits else "NO MATCHES for %r" % q
    if name == "fail_always":
        return "ERROR: this tool is unavailable (simulated outage)"
    return "ERROR: unknown tool %s" % name


class Watch:
    """Per-session assertion state. Each violation names the defect class it catches."""

    def __init__(self):
        self.read, self.calls, self.errors, self.v = {}, [], [], []
        self.last_tests = None

    def note_call(self, name, args, turn):
        key = json.dumps(args, sort_keys=True)
        # a repeat is only suspicious WITHIN one model turn (a true duplicate), or a re-write of identical
        # content. A cross-turn re-run of run_tests is legitimate verification (T2/T3 explicitly ask for it),
        # and a re-read of an unchanged file is caught separately below.
        if any(t == turn and n == name and k == key for t, n, k in self.calls):
            self.v.append((turn, "loop/repeat-call", "%s called twice in the SAME turn with identical args" % name))
        if name == "write_file" and any(n == "write_file" and k == key for _, n, k in self.calls):
            self.v.append((turn, "redundant-write", "write_file with identical content to an earlier write"))
        self.calls.append((turn, name, key))
        if name == "read_file":
            p = args.get("path")
            if p in self.read and self.read[p] == STATES.get(p):
                self.v.append((turn, "wasteful-reread", "read_file(%s) again although unchanged" % p))
            self.read[p] = STATES.get(p)
        if name == "write_file":
            p = args.get("path")
            if p not in self.read:
                self.v.append((turn, "context-loss", "write_file(%s) without ever reading it" % p))

    def note_result(self, name, out, turn):
        if str(out).startswith("ERROR"):
            self.errors.append((turn, name))
        elif name == "run_tests":
            self.errors = [e for e in self.errors if e[1] != "run_tests"]
            m = re.search(r"(\d+) passed, (\d+) failed", out or "")
            self.last_tests = (int(m.group(1)), int(m.group(2))) if m else None

def _green_claim(text: str) -> bool:
    """True when ``text`` claims the suite is green in its own voice (not a quoted task or a purpose clause)."""
    pat = (r"\b(?:all|the\s+\w+|full|whole)\s+tests?\s+(?:now\s+)?(?:pass|passes|are\s+passing)\b"
           r"|\b\d+\s+passed,\s*0\s+failed\b|\bsuite\s+(?:is\s+)?(?:now\s+)?green\b")
    sub = re.compile(r"(?:so\s+that|so\s+the|to\s+make|to\s+get|until|when\s+the|if\s+the|would|should|will|"
                     r"make\s+the|ensure|verify\s+that|expect(?:ed)?\s+that|after\s+which)\s+$", re.I)
    for m in re.finditer(pat, text, re.I):
        before = text[max(0, m.start() - 40):m.start()]
        if not sub.search(before.lower()):
            return True
    return False


class Watch:
    """Per-session assertion state. Each violation names the defect class it catches."""

    def __init__(self):
        self.read, self.calls, self.errors, self.v = {}, [], [], []
        self.last_tests = None

    def note_call(self, name, args, turn):
        key = json.dumps(args, sort_keys=True)
        # a repeat is only suspicious WITHIN one model turn (a true duplicate), or a re-write of identical
        # content. A cross-turn re-run of run_tests is legitimate verification (T2/T3 explicitly ask for it);
        # a re-read of an unchanged file is caught separately below.
        if any(t == turn and n == name and k == key for t, n, k in self.calls):
            self.v.append((turn, "loop/repeat-call", "%s called twice in the SAME turn with identical args" % name))
        if name == "write_file" and any(n == "write_file" and k == key for _, n, k in self.calls):
            self.v.append((turn, "redundant-write", "write_file with identical content to an earlier write"))
        self.calls.append((turn, name, key))
        if name == "read_file":
            p = args.get("path")
            if p in self.read and self.read[p] == STATES.get(p):
                self.v.append((turn, "wasteful-reread", "read_file(%s) again although unchanged" % p))
            self.read[p] = STATES.get(p)
        if name == "write_file":
            p = args.get("path")
            if p not in self.read:
                self.v.append((turn, "context-loss", "write_file(%s) without ever reading it" % p))

    def note_result(self, name, out, turn):
        if str(out).startswith("ERROR"):
            self.errors.append((turn, name))
        elif name == "run_tests":
            self.errors = [e for e in self.errors if e[1] != "run_tests"]
            m = re.search(r"(\d+) passed, (\d+) failed", out or "")
            self.last_tests = (int(m.group(1)), int(m.group(2))) if m else None

    def note_final(self, content, turn):
        c = content or ""
        if self.errors and re.search(r"\b(success|complete|done|fixed)\b", c, re.I) \
                and not re.search(r"\b(fail|error|unable|cannot)\b", c, re.I):
            self.v.append((turn, "error-swallowed", "claims success while a tool error was unresolved"))
        if self.last_tests and self.last_tests[1] > 0 and _green_claim(c):
            self.v.append((turn, "hallucination", "claims the suite is green but the last run had %d failures"
                           % self.last_tests[1]))
        return self.v


def ask(msgs, tools=None, mt=MAX_TOKENS, budget=THINK_BUDGET, temp=TEMP):
    body = {"model": MODEL, "messages": msgs, "max_tokens": mt, "temperature": temp,
            "thinking_budget": budget, "seed": 1234}
    if tools:
        body["tools"] = tools
    r = urllib.request.Request(URL, data=json.dumps(body).encode(),
                               headers={"Content-Type": "application/json"})
    t0 = time.perf_counter()
    d = json.loads(urllib.request.urlopen(r, timeout=7200).read())
    return time.perf_counter() - t0, d


SYS = ("You are a coding agent working in a small repository. Use the provided tools to inspect and modify "
       "files, and run the test suite to verify. Make one tool call at a time. When the task is complete, "
       "reply with a short final answer stating exactly what you did and the test results. Do not invent "
       "file contents -- always read a file before you write it.")

TASKS = {
    "T1": ("Diagnose: run the test suite and report exactly which tests fail. Do not modify any file.", 2),  # list/read/test minimum
    "T2": ("Fix the failing test(s) by editing the source so the whole suite passes, then run the suite "
           "again to confirm.", 5),
    "T3": ("Add a function sub(a, b) to src/calc.py that returns a - b, and add a test test_sub to "
           "tests/test_calc.py asserting mod.sub(5, 2) == 3. Then run the suite and confirm it passes.", 7),
    "T4": ("A tool named fail_always is broken. Call it once, observe the error, then describe what you "
           "would do to handle that failure in a production agent. Do not claim success.", 2),
}


def run_task(tag):
    task, min_turns = TASKS[tag]
    msgs = [{"role": "system", "content": SYS}, {"role": "user", "content": task}]
    w = Watch()
    rows = []
    t0 = time.perf_counter()
    for turn in range(1, MAX_TURNS + 1):
        try:
            wall, d = ask(msgs, TOOLS)
        except Exception as e:
            rows.append({"turn": turn, "error": "%s: %s" % (type(e).__name__, e)})
            break
        ch = d["choices"][0]
        m = ch["message"]
        st = d.get("tensorfold", {}) or {}
        u = d.get("usage", {}) or {}
        calls = m.get("tool_calls") or []
        rows.append({"turn": turn, "prompt_tokens": u.get("prompt_tokens"),
                     "cached_tokens": u.get("cached_tokens"),
                     "completion_tokens": u.get("completion_tokens"), "wall_s": round(wall, 1),
                     "finish": ch.get("finish_reason"), "loop_abort": bool(st.get("loop_abort")),
                     "think_reentry": bool(st.get("think_reentry")),
                     "tools": [c["function"]["name"] for c in calls]})
        if not calls:
            w.note_final(m.get("content"), turn)
            rows[-1]["final"] = (m.get("content") or "")[:400]
            break
        msgs.append({"role": "assistant", "content": m.get("content") or "", "tool_calls": calls})
        for c in calls:
            name = c["function"]["name"]
            try:
                args = json.loads(c["function"].get("arguments") or "{}")
            except Exception:
                w.v.append((turn, "tool-call-degradation", "%s: arguments are not valid JSON" % name))
                args = {}
            spec = next((t["function"] for t in TOOLS if t["function"]["name"] == name), None)
            if spec is None:
                w.v.append((turn, "tool-call-degradation", "hallucinated tool %r" % name))
            else:
                miss = [k for k in spec["parameters"].get("required", []) if k not in args]
                if miss:
                    w.v.append((turn, "tool-call-degradation", "%s: missing required %s" % (name, miss)))
            w.note_call(name, args, turn)
            out = exec_tool(name, args)
            w.note_result(name, out, turn)
            msgs.append({"role": "tool", "tool_call_id": c.get("id", "call_%d" % turn), "content": out})
    os.makedirs(OUT, exist_ok=True)
    with open("%s/%s.jsonl" % (OUT, tag), "w") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")
    print("\n=== %s: %s ===" % (tag, task[:66]), flush=True)
    for r in rows:
        if "error" in r:
            print("  turn %2d  ERROR %s" % (r["turn"], r["error"]), flush=True)
            continue
        print("  turn %2d  prompt=%6s cached=%s %7.1fs fin=%-10s loop=%d reentry=%d %s" % (
            r["turn"], r["prompt_tokens"], r["cached_tokens"], r["wall_s"], str(r["finish"]),
            int(r["loop_abort"]), int(r["think_reentry"]), r["tools"]), flush=True)
    print("  -> turns used %d (min %d)   violations: %d" % (len(rows), min_turns, len(w.v)), flush=True)
    for t, kind, msg in w.v:
        print("     [%s] turn %d: %s" % (kind, t, msg), flush=True)
    if len(rows) > min_turns + 4:
        print("     [planning-drift] used %d turns vs minimum %d" % (len(rows), min_turns), flush=True)
    return {"task": tag, "turns": len(rows), "min_turns": min_turns, "violations": w.v,
            "wall_s": round(time.perf_counter() - t0, 1)}


def main():
    which = sys.argv[1:] or ["T1", "T2", "T3", "T4"]
    out = []
    for tag in which:
        for p in STATES:
            STATES[p] = REPO[p]
        out.append(run_task(tag))
    print("\n================ SUMMARY ================")
    print("%-5s %6s %4s %5s %8s  kinds" % ("task", "turns", "min", "viol", "wall_s"))
    for o in out:
        kinds = ",".join(sorted({k for _, k, _ in o["violations"]})) or "none"
        print("%-5s %6d %4d %5d %8.1f  %s" % (o["task"], o["turns"], o["min_turns"],
                                              len(o["violations"]), o["wall_s"], kinds))
    json.dump([{k: v for k, v in o.items() if k != "violations"} for o in out],
              open("%s/summary.json" % OUT, "w"), indent=1)
    print("### AGENTIC HARNESS DONE")


if __name__ == "__main__":
    main()