#!/usr/bin/env python3
"""chipdesk: a timing-repair agent that runs real OpenROAD on a checkpoint.

  python3 chipdesk.py triage <report_checks.rpt>
  python3 chipdesk.py fix --ckpt 4_cts.odb --sdc 4_cts.sdc [--llm ollama:qwen2.5:7b | workers | none] [--trace out.json]
  python3 chipdesk.py selftest

The LLM only proposes. Every proposal is checked against the command registry parsed from
OpenROAD/OpenSTA define_cmd_args, restricted to an allowlist, run in OpenROAD on a copy of the
checkpoint, re-measured, and accepted only if the assertions pass. Standard library only.
"""
import argparse, hashlib, json, math, os, re, shlex, subprocess, sys, time, urllib.request
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent
CORPUS = ROOT / "app" / "public" / "corpus.json"
IMAGE = "openroad/orfs:latest"
OPENROAD = "/OpenROAD-flow-scripts/tools/install/OpenROAD/bin/openroad"
PLATFORM = "/OpenROAD-flow-scripts/flow/platforms/nangate45"

# Steps the agent may run. Each flag must also be documented for the command in the registry.
ALLOW = {
    "repair_timing": {"-setup", "-hold", "-setup_margin", "-hold_margin", "-repair_tns", "-max_passes",
                      "-sequence", "-skip_pin_swap", "-skip_gate_cloning", "-skip_buffering", "-allow_setup_violations", "-verbose"},
    "repair_design": {"-max_wire_length", "-slew_margin", "-cap_margin"},
}
BOUNDS = {"-setup_margin": (0, 0.2), "-hold_margin": (0, 0.2), "-repair_tns": (0, 100),
          "-max_passes": (1, 100000), "-max_wire_length": (1, 5000), "-slew_margin": (0, 50), "-cap_margin": (0, 50)}
SEQ_MOVES = {"unbuffer", "vt_swap", "sizeup", "size_down", "swap", "buffer", "clone", "split"}


# ---------------------------------------------------------------- registry + docs
def load_corpus():
    return json.loads(CORPUS.read_text())


def tokenize(s):
    stop = set("a an the and or of to in on for with by is are it this that as at from be can will".split())
    return [t.rstrip(".") for t in re.findall(r"-?[a-z0-9_][a-z0-9_.]*", s.lower()) if t not in stop and len(t) > 1]


def bm25(chunks, query, k=4, k1=1.2, b=0.75):
    docs = [tokenize(c["heading"] + " " + c["text"]) for c in chunks]
    df = {}
    for d in docs:
        for t in set(d):
            df[t] = df.get(t, 0) + 1
    n, avg = len(docs), sum(map(len, docs)) / len(docs)
    q = set(tokenize(query))
    scores = []
    for i, d in enumerate(docs):
        tf, s = {}, 0.0
        for t in d:
            tf[t] = tf.get(t, 0) + 1
        for t in q:
            f = tf.get(t)
            if f:
                idf = math.log(1 + (n - df[t] + 0.5) / (df[t] + 0.5))
                s += idf * f * (k1 + 1) / (f + k1 * (1 - b + b * len(d) / avg))
        if s:
            scores.append((s, i))
    return [chunks[i] for _, i in sorted(scores, reverse=True)[:k]]


# ---------------------------------------------------------------- report parser
def parse_timing(text):
    """Parse OpenSTA report_checks. Statistics cover the reported paths only, not the design."""
    paths = []
    for blk in re.split(r"^\s*Startpoint:", text, flags=re.M)[1:]:
        blk = "Startpoint:" + blk
        sl = re.search(r"(-?\d+\.?\d*)\s+slack\s*\((MET|VIOLATED)\)", blk)
        if not sl:
            continue
        g = lambda p: (re.search(p, blk) or [None, None])[1]
        stages = []
        for line in blk.split("data arrival time")[0].splitlines():
            m = re.match(r"^\s*((?:-?\d+\.?\d*\s+)+)([\^v])\s+(\S+)\s+\(([^)]+)\)", line)
            if m:
                nums = [float(x) for x in m.group(1).split()]
                if len(nums) >= 2:
                    stages.append({"incr": nums[-2], "time": nums[-1], "pin": m.group(3), "cell": m.group(4)})
        worst = max(stages, key=lambda s: s["incr"]) if stages else None
        paths.append({"start": g(r"Startpoint:\s*(\S+)"), "end": g(r"Endpoint:\s*(\S+)"),
                      "type": "hold" if g(r"Path Type:\s*(\w+)") == "min" else "setup",
                      "slack": float(sl.group(1)), "stages": len(stages), "worst_stage": worst})
    def summ(kind):
        ps = [p for p in paths if p["type"] == kind]
        return {"reported_paths": len(ps), "violating": sum(p["slack"] < 0 for p in ps),
                "worst_reported_slack": min((p["slack"] for p in ps), default=None)}
    return {"paths": paths, "setup": summ("setup"), "hold": summ("hold"),
            "note": "covers reported paths only; design WNS/TNS come from report_wns/report_tns"}


# ---------------------------------------------------------------- verifier
def verify_step(tcl, registry):
    """Returns (ok, issues). One or two lines, allowlisted commands, documented flags, legal values."""
    issues = []
    lines = [l.strip() for l in tcl.strip().splitlines() if l.strip() and not l.strip().startswith("#")]
    if not 1 <= len(lines) <= 2:
        return False, [f"expected 1 or 2 command lines, got {len(lines)}"]
    for line in lines:
        if re.search(r"[;$\[\]`]|\bexec\b|\bsource\b|\bexit\b", line):
            issues.append(f"forbidden Tcl syntax in: {line}")
            continue
        try:
            toks = shlex.split(line)
        except ValueError as e:
            issues.append(f"cannot parse: {e}")
            continue
        cmd, rest = toks[0], toks[1:]
        if cmd not in registry:
            issues.append(f"{cmd}: not an OpenROAD/OpenSTA command")
            continue
        if cmd not in ALLOW:
            issues.append(f"{cmd}: real command, but not on the allowlist for this checkpoint")
            continue
        entry = registry[cmd]
        arity = entry.get("arity", {})
        i = 0
        while i < len(rest):
            f = rest[i]
            if not f.startswith("-"):
                issues.append(f"{cmd}: unexpected positional '{f}'")
                i += 1
                continue
            if f not in entry["flags"]:
                issues.append(f"{cmd} {f}: flag not documented for {cmd}")
            elif f not in ALLOW[cmd]:
                issues.append(f"{cmd} {f}: documented, but not allowlisted")
            takes = arity.get(f, False)
            if takes:
                if i + 1 >= len(rest) or rest[i + 1].startswith("-") and not re.match(r"^-?\d", rest[i + 1]):
                    issues.append(f"{cmd} {f}: needs a value")
                else:
                    v = rest[i + 1]
                    if f in BOUNDS:
                        try:
                            x = float(v)
                            lo, hi = BOUNDS[f]
                            if not lo <= x <= hi:
                                issues.append(f"{cmd} {f} {v}: outside {lo}..{hi}")
                        except ValueError:
                            issues.append(f"{cmd} {f} {v}: not a number")
                    if f == "-sequence" and not set(v.replace(",", " ").split()) <= SEQ_MOVES:
                        issues.append(f"{cmd} -sequence: unknown move in '{v}'")
                    i += 1
            i += 1
        if cmd == "repair_timing" and not ({"-setup", "-hold"} & set(rest)):
            issues.append("repair_timing: say -setup or -hold, so the assertion knows what to check")
        if {"-setup", "-hold"} <= set(rest):
            issues.append("repair_timing: one class per step (-setup or -hold), so each change is checked on its own")
    return not issues, issues


# ---------------------------------------------------------------- OpenROAD runner
def run_openroad(ckpt, sdc, step_tcl=None, timeout=600):
    work = ROOT
    rel = lambda p: "/work/" + str(Path(p).resolve().relative_to(work))
    env = f"PLATFORM_DIR={PLATFORM} CKPT={rel(ckpt)} SDC={rel(sdc)}"
    if step_tcl is not None:
        sf = HERE / "tcl" / "_step.tcl"
        sf.write_text(step_tcl + "\n")
        env += f" STEP_TCL={rel(sf)}"
    cmd = ["docker", "run", "--rm", "--network", "none", "-u", f"{os.getuid()}:{os.getgid()}",
           "-v", f"{work}:/work", IMAGE, "bash", "-c",
           f"export {env}; {OPENROAD} -no_init -exit /work/cli/tcl/measure.tcl"]
    t0 = time.time()
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        out = p.stdout + p.stderr
    except subprocess.TimeoutExpired:
        return {"ok": False, "error": f"timeout after {timeout}s", "seconds": timeout}
    metrics = {}
    for tag, key, val in re.findall(r"^METRIC (\w+) (\w+) (\S+)", out, re.M):
        metrics.setdefault(tag, {})[key] = float(val)
    err = re.search(r"^STEP_ERROR (.*)$", out, re.M) or re.search(r"^\[ERROR [^\]]+\].*$", out, re.M)
    return {"ok": p.returncode == 0 and not err, "error": err.group(0) if err else None,
            "metrics": metrics, "seconds": round(time.time() - t0, 1),
            "report_tail": "\n".join(l for l in out.splitlines() if "slack (" in l)[:400],
            "log_tail": "\n".join(out.splitlines()[-12:])}


def assertions(before, after, kind, base_area=None, budget=5.0):
    """A step counts as a fix only if all of these hold. Area is measured against the original checkpoint."""
    checks = []
    if kind == "setup":
        checks.append(("setup TNS improves by at least 2%", after["setup_tns"] > before["setup_tns"] * 0.98 if before["setup_tns"] < 0 else True))
        checks.append(("setup WNS does not get worse", after["setup_wns"] >= before["setup_wns"] - 1e-4))
        checks.append(("hold WNS stays >= min(0, before)", after["hold_wns"] >= min(0.0, before["hold_wns"]) - 1e-4))
    else:
        checks.append(("hold WNS improves", after["hold_wns"] > before["hold_wns"]))
        checks.append(("setup TNS gets no more than 1% worse", after["setup_tns"] >= before["setup_tns"] * 1.01 - 1e-4))
    base = base_area or before["area"]
    checks.append((f"area within {budget:g}% of the original", after["area"] <= base * (1 + budget / 100) + 1e-9))
    return [{"check": c, "pass": bool(ok)} for c, ok in checks]


# ---------------------------------------------------------------- LLM backends
def llm(backend, messages):
    if backend == "none":
        return None, "none"
    if backend.startswith("ollama:"):
        model = backend.split(":", 1)[1]
        host = os.environ.get("OLLAMA_HOST", "127.0.0.1:11434")
        req = urllib.request.Request(f"http://{host}/api/chat", data=json.dumps(
            {"model": model, "messages": messages, "stream": False, "format": "json",
             "options": {"temperature": 0}}).encode(), headers={"content-type": "application/json"})
        with urllib.request.urlopen(req, timeout=600) as r:
            return json.loads(r.read())["message"]["content"], f"ollama/{model} (local)"
    if backend == "workers":
        req = urllib.request.Request("https://oppstar-agent.pages.dev/api/llm", data=json.dumps(
            {"messages": messages, "json": True}).encode(), headers={"content-type": "application/json", "user-agent": "chipdesk"})
        with urllib.request.urlopen(req, timeout=120) as r:
            d = json.loads(r.read())
            return d["text"], d.get("model", "workers-ai")
    raise SystemExit(f"unknown --llm {backend}")


RULE_STEPS = {"hold": ["repair_timing -hold", "repair_timing -hold -hold_margin 0.02"],
              "setup": ["repair_timing -setup", "repair_timing -setup -sequence unbuffer,sizeup,swap,buffer -repair_tns 100", "repair_design"]}


def propose(backend, state, registry, docs, feedback, attempt):
    if backend == "none":
        # hold first: it is cheap and rarely conflicts; then setup moves in order
        queue = [t for k in ("hold", "setup") if k in state["open_violations"] for t in RULE_STEPS[k]]
        done = {f["tcl"] for f in feedback} | set(state["accepted_steps"])
        nxt = next((t for t in queue if t not in done), queue[-1] if queue else "repair_design")
        return {"tcl": nxt, "why": "rule fallback"}, "rules"
    allow = {c: {f: ("value" if registry[c].get("arity", {}).get(f) else "switch") for f in sorted(ALLOW[c]) if f in registry[c]["flags"]} for c in ALLOW}
    sys_msg = ("You are a timing-closure agent for OpenROAD. Propose ONE step to run on a post-CTS checkpoint. "
               "Use only the allowlisted commands and flags below. Give values for flags that take one. "
               'Reply as JSON: {"tcl": "<one command line>", "why": "<one sentence citing a doc passage [n] or a metric>"}.\n'
               f"ALLOWLIST: {json.dumps(allow)}\nBOUNDS: {json.dumps(BOUNDS)}\n"
               "DOCS:\n" + "\n\n".join(f"[{i+1}] {d['title']} > {d['heading']}\n{d['text'][:900]}" for i, d in enumerate(docs)))
    user = f"STATE: {json.dumps(state)}"
    if feedback:
        user += f"\nPREVIOUS ATTEMPTS AND RESULTS: {json.dumps(feedback)}\nDo not repeat a step that failed. Try a different documented option."
    text, model = llm(backend, [{"role": "system", "content": sys_msg}, {"role": "user", "content": user}])
    try:
        m = re.search(r"\{.*\}", text, re.S)
        d = json.loads(m.group(0))
        return {"tcl": str(d.get("tcl", "")).strip(), "why": str(d.get("why", ""))[:300]}, model
    except Exception:
        return {"tcl": "", "why": f"unparseable LLM output: {text[:200]}"}, model


# ---------------------------------------------------------------- commands
def sha(p):
    return hashlib.sha256(Path(p).read_bytes()).hexdigest()[:16]


def image_digest():
    try:
        return subprocess.run(["docker", "image", "inspect", IMAGE, "--format", "{{index .RepoDigests 0}}"],
                              capture_output=True, text=True, timeout=30).stdout.strip()
    except Exception:
        return "unknown"


def cmd_fix(a):
    corpus = load_corpus()
    reg = corpus["registry"]
    docs = bm25(corpus["chunks"], "repair_timing setup violations margin sequence repair_tns post CTS propagated clock", k=4)
    trace = {"tool": "chipdesk fix", "started": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
             "inputs": {"ckpt": str(a.ckpt), "ckpt_sha256": sha(a.ckpt), "sdc": str(a.sdc), "sdc_sha256": sha(a.sdc)},
             "versions": {"openroad_image": image_digest(), "corpus_openroad_sha": corpus["source"]["sha"],
                          "registry_opensta_sha": corpus["source"].get("sta", {}).get("sha"), "llm_backend": a.llm},
             "docs_used": [d["id"] for d in docs], "steps": []}
    base = run_openroad(a.ckpt, a.sdc)
    if not base["ok"]:
        raise SystemExit(f"baseline failed: {base['error']}\n{base['log_tail']}")
    before = base["metrics"]["before"]
    trace["baseline"] = before
    print(f"baseline  setup WNS {before['setup_wns']:+.4f}  TNS {before['setup_tns']:+.4f}  hold WNS {before['hold_wns']:+.4f}  area {before['area']:.1f}")
    open_classes = [k for k in ("setup", "hold") if before[f"{k}_wns"] < 0]
    if not open_classes:
        print("no violations; nothing to do")
        trace["status"] = "clean"
        return finish(trace, a)
    accepted, current, feedback, tried = [], before, [], set()
    for attempt in range(a.attempts):
        if not open_classes:
            break
        t0 = time.time()
        state = {"open_violations": open_classes, "metrics_now": current, "accepted_steps": accepted}
        prop, model = propose(a.llm, state, reg, docs, feedback, attempt)
        ok, issues = verify_step(prop["tcl"], reg) if prop["tcl"] else (False, [prop["why"]])
        norm = " ".join(prop["tcl"].split())
        if ok and norm in tried:
            ok, issues = False, ["exact repeat of an earlier attempt"]
        step = {"attempt": attempt + 1, "model": model, "proposal": prop, "verifier": {"ok": ok, "issues": issues},
                "llm_seconds": round(time.time() - t0, 1)}
        print(f"\n#{attempt+1} {model}: {prop['tcl'] or '(none)'}\n   why: {prop['why']}")
        if not ok:
            print("   verifier REJECTED: " + "; ".join(issues))
            feedback.append({"tcl": prop["tcl"], "rejected_by_verifier": issues})
            trace["steps"].append(step)
            continue
        kind = "hold" if re.search(r"(^|\s)-hold(\s|$)", prop["tcl"]) else "setup"
        if kind not in open_classes:
            print(f"   REJECTED: targets {kind}, which has no violation")
            step["verifier"] = {"ok": False, "issues": [f"targets {kind}, which has no violation"]}
            feedback.append({"tcl": prop["tcl"], "rejected": f"{kind} is not violated; open: {open_classes}"})
            trace["steps"].append(step)
            continue
        tried.add(norm)
        run = run_openroad(a.ckpt, a.sdc, "\n".join(accepted + [prop["tcl"]]))
        step["run"] = {k: run[k] for k in ("ok", "error", "seconds")}
        if not run["ok"]:
            print(f"   OpenROAD ERROR: {run['error']}")
            feedback.append({"tcl": prop["tcl"], "openroad_error": run["error"]})
            trace["steps"].append(step)
            continue
        after = run["metrics"]["after"]
        checks = assertions(current, after, kind, base_area=before["area"], budget=a.area_budget)
        step.update(target=kind, after=after, assertions=checks)
        passed = all(c["pass"] for c in checks)
        print(f"   after   setup WNS {after['setup_wns']:+.4f}  TNS {after['setup_tns']:+.4f}  hold WNS {after['hold_wns']:+.4f}  area {after['area']:.1f}  ({run['seconds']}s)")
        for c in checks:
            print(f"   {'PASS' if c['pass'] else 'FAIL'}  {c['check']}")
        trace["steps"].append(step)
        if passed:
            accepted.append(prop["tcl"])
            current = after
            step["accepted"] = True
            open_classes = [k for k in ("setup", "hold") if current[f"{k}_wns"] < 0]
            print(f"   ACCEPTED. open violations now: {open_classes or 'none'}")
        else:
            feedback.append({"tcl": prop["tcl"], "target": kind, "after": after,
                             "failed_checks": [c["check"] for c in checks if not c["pass"]]})
    trace["accepted"] = accepted
    trace["final"] = current
    trace["classes"] = {k: ("closed" if current[f"{k}_wns"] >= 0 else "open") for k in ("setup", "hold")}
    trace["status"] = "fixed" if not open_classes else ("partial" if accepted else "not_fixed")
    if "setup" in open_classes:
        # a deterministic diagnosis the engineer can act on
        trace["diagnosis"] = min_period_note(before)
    print(f"\nstatus: {trace['status']}  classes: {trace['classes']}  accepted: {accepted or 'none'}")
    if trace.get("diagnosis"):
        print("diagnosis: " + trace["diagnosis"])
    return finish(trace, a)


def min_period_note(m):
    return (f"Setup WNS {m['setup_wns']:+.4f} ns with no allowlisted step meeting the assertions. "
            "Sizing and buffering moves are exhausted at this clock; the next levers are outside this agent's "
            "allowlist: relax the clock period, restructure logic (rmp), or re-synthesise with a tighter target. "
            "Escalate to the PD engineer.")


def finish(trace, a):
    trace["finished"] = time.strftime("%Y-%m-%dT%H:%M:%S%z")
    if a.trace:
        Path(a.trace).write_text(json.dumps(trace, indent=1))
        print(f"trace -> {a.trace}")
    return 0


def cmd_triage(a):
    text = Path(a.report).read_text()
    p = parse_timing(text)
    tool = {}
    for key, pat in (("setup_tns", r"^tns max (\S+)"), ("setup_wns", r"^wns max (\S+)"), ("worst_slack_max", r"^worst slack max (\S+)"),
                     ("worst_slack_min", r"^worst slack min (\S+)")):
        m = re.search(pat, text, re.M)
        if m:
            tool[key] = float(m.group(1))
    out = {"parsed": {k: p[k] for k in ("setup", "hold", "note")}, "tool_native": tool,
           "worst_paths": [x for x in p["paths"] if x["slack"] < 0][:3]}
    if "worst_slack_max" in tool and p["setup"]["worst_reported_slack"] is not None:
        out["cross_check"] = {"parsed_worst_setup": p["setup"]["worst_reported_slack"], "tool_worst_setup": tool["worst_slack_max"],
                              "match": abs(p["setup"]["worst_reported_slack"] - tool["worst_slack_max"]) < 0.006}
    print(json.dumps(out, indent=1))
    return 0


def cmd_selftest(_):
    corpus = load_corpus()
    reg = corpus["registry"]
    n = 0
    def t(cond, msg):
        nonlocal n
        assert cond, msg
        n += 1
    t(verify_step("repair_timing -setup -setup_margin 0.05", reg)[0], "clean step passes")
    t(not verify_step("repair_timing -setup -setup_margin banana", reg)[0], "non-number value rejected")
    t(not verify_step("repair_timing -setup -setup_margin", reg)[0], "missing value rejected")
    t(not verify_step("repair_timing -setup -fix_all_paths", reg)[0], "invented flag rejected")
    t(not verify_step("fix_hold_violations -margin 0.05", reg)[0], "invented command rejected")
    t(not verify_step("detailed_route", reg)[0], "real but not allowlisted rejected")
    t(not verify_step("repair_timing -setup; exec rm -rf /", reg)[0], "Tcl injection rejected")
    t(not verify_step("repair_timing -setup -setup_margin 5", reg)[0], "out of bounds rejected")
    t(not verify_step("repair_timing -setup -sequence sizeup,teleport", reg)[0], "unknown move rejected")
    t(not verify_step("repair_timing", reg)[0], "no -setup/-hold rejected")
    t(not verify_step("repair_timing -setup -hold", reg)[0], "both classes in one step rejected")
    t("report_checks" in reg and "set_propagated_clock" in reg, "OpenSTA commands present")
    fx = (ROOT / "app" / "public" / "fixtures" / "gcd_4_cts_final.rpt")
    if fx.exists():
        p = parse_timing(fx.read_text())
        t(p["setup"]["violating"] >= 1, "real report parsed")
    print(f"{n} selftests passed")
    return 0


def main():
    ap = argparse.ArgumentParser(prog="chipdesk", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("triage"); s.add_argument("report"); s.set_defaults(fn=cmd_triage)
    s = sub.add_parser("fix"); s.add_argument("--ckpt", required=True); s.add_argument("--sdc", required=True)
    s.add_argument("--llm", default="ollama:qwen2.5:7b"); s.add_argument("--attempts", type=int, default=4)
    s.add_argument("--area-budget", type=float, default=5.0)
    s.add_argument("--trace"); s.set_defaults(fn=cmd_fix)
    s = sub.add_parser("selftest"); s.set_defaults(fn=cmd_selftest)
    a = ap.parse_args()
    sys.exit(a.fn(a))


if __name__ == "__main__":
    main()