#!/usr/bin/env python3
"""
tree.py -- persistent hypothesis-tree state manager for Arbor-style
Hypothesis Tree Refinement (HTR).

The hypothesis tree is the durable research state for an Autonomous
Optimization (AO) run. This script owns the *mechanical* parts of that state
-- creating nodes, writing back evidence, propagating insights up the tree,
pruning falsified branches, recording the held-out merge gate, and rendering
an "Observe" projection -- so the coordinator (you, the model) can spend its
judgment on what the evidence *means* rather than on bookkeeping.

Division of labor:
  - This script keeps the state consistent and auditable. It never decides
    which hypothesis to try or whether one is good.
  - The coordinator reads the projection (`observe`), forms hypotheses,
    interprets executor reports, and calls the mutating commands to record
    those decisions.

State lives in `.arbor/` under the run directory (default: current dir):
  .arbor/tree.json   -- the hypothesis tree (nodes, edges, evidence, insights)
  .arbor/run.json    -- run-level config: objective, evaluators, budget, M_best

Node fields mirror the paper's research unit  n = <h, iota, mu>:
  hypothesis  (h)    -- the falsifiable claim this node tests
  insight     (iota) -- distilled, reusable lesson (filled after execution)
  metadata    (mu)   -- status, dev_score, test_score, result, branch_ref, depth

Run `python tree.py --help` or `python tree.py <command> --help`.
"""

import argparse
import json
import math
import sys
import time
from pathlib import Path

# ----------------------------------------------------------------------------
# Storage helpers
# ----------------------------------------------------------------------------

VALID_STATUS = {"pending", "running", "executed", "merged", "pruned", "root"}


def _finite_score(value):
    """Reject failed-evaluation sentinels instead of persisting invalid scores."""
    score = float(value)
    if not math.isfinite(score):
        raise argparse.ArgumentTypeError("score must be finite (not NaN or infinity)")
    return score


def _dir(run_dir):
    return Path(run_dir) / ".arbor"


def _tree_path(run_dir):
    return _dir(run_dir) / "tree.json"


def _run_path(run_dir):
    return _dir(run_dir) / "run.json"


def _load(path, what):
    if not path.exists():
        sys.exit(
            f"error: no {what} found at {path}. Run `tree.py init` first "
            f"(from the run directory, or pass --run-dir)."
        )
    with open(path) as f:
        return json.load(f)


def _save(path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    with open(tmp, "w") as f:
        json.dump(data, f, indent=2)
    tmp.replace(path)


def _load_tree(run_dir):
    return _load(_tree_path(run_dir), "hypothesis tree")


def _load_run(run_dir):
    return _load(_run_path(run_dir), "run config")


def _next_id(tree):
    n = tree.get("_counter", 0) + 1
    tree["_counter"] = n
    return f"n{n}"


def _id_sort_key(node_id):
    """Order generated ids numerically, then any manually assigned ids by name."""
    if isinstance(node_id, str) and node_id.startswith("n") and node_id[1:].isdecimal():
        return (0, int(node_id[1:]), node_id)
    return (1, 0, str(node_id))


def _node(tree, node_id):
    node = tree["nodes"].get(node_id)
    if node is None:
        sys.exit(f"error: no node with id '{node_id}'. Run `tree.py status` to list nodes.")
    return node


def _children(tree, node_id):
    return [nid for nid, n in tree["nodes"].items() if n.get("parent") == node_id]


def _ancestors(tree, node_id):
    """Path from node up to (and including) root, nearest first."""
    path = []
    cur = tree["nodes"][node_id].get("parent")
    while cur is not None:
        path.append(cur)
        cur = tree["nodes"][cur].get("parent")
    return path


def _depth(tree, node_id):
    return len(_ancestors(tree, node_id))


def _stamp(node):
    node["updated_at"] = int(time.time())


# ----------------------------------------------------------------------------
# Commands
# ----------------------------------------------------------------------------


def cmd_init(args):
    d = _dir(args.run_dir)
    if _tree_path(args.run_dir).exists() and not args.force:
        sys.exit(
            f"error: a tree already exists at {_tree_path(args.run_dir)}. "
            f"Use --force to overwrite (this erases the current run)."
        )
    root = {
        "id": "n0",
        "parent": None,
        "depth": 0,
        "status": "root",
        "hypothesis": args.objective,
        "insight": "",  # global insights accumulate here via backpropagation
        "metadata": {"dev_score": None, "test_score": None, "result": "", "branch_ref": None},
        "created_at": int(time.time()),
        "updated_at": int(time.time()),
    }
    tree = {"_counter": 0, "root": "n0", "nodes": {"n0": root}}
    run = {
        "objective": args.objective,
        "metric_direction": args.metric_direction,
        "dev_eval": args.dev_eval,
        "test_eval": args.test_eval,
        "material": args.material,
        "branching": args.branching,
        "max_depth": args.max_depth,
        "budget_cycles": args.budget,
        "cycles_used": 0,
        "best_node": None,           # node id of current M_best
        "best_test_score": None,     # held-out score of M_best
        "best_branch_ref": None,     # git ref / path of M_best artifact
        "created_at": int(time.time()),
    }
    _save(_tree_path(args.run_dir), tree)
    _save(_run_path(args.run_dir), run)
    print(f"Initialized Arbor run in {d}")
    print(f"  objective       : {args.objective}")
    print(f"  metric direction: {args.metric_direction} (higher-is-better after orienting)")
    print(f"  dev evaluator   : {args.dev_eval}")
    print(f"  test evaluator  : {args.test_eval}")
    print(f"  budget          : {args.budget} cycles, branching {args.branching}, max depth {args.max_depth}")
    print("\nNext: `tree.py observe` to read the state, then add direction nodes under n0.")


def cmd_add_node(args):
    tree = _load_tree(args.run_dir)
    parent = _node(tree, args.parent)
    run = _load_run(args.run_dir)
    nid = _next_id(tree)
    depth = _depth(tree, args.parent) + 1
    if depth > run["max_depth"]:
        print(
            f"warning: node depth {depth} exceeds max_depth {run['max_depth']}. "
            f"Deep nodes should be concrete, executable interventions.",
            file=sys.stderr,
        )
    node = {
        "id": nid,
        "parent": args.parent,
        "depth": depth,
        "status": "pending",
        "hypothesis": args.hypothesis,
        "insight": "",
        "metadata": {"dev_score": None, "test_score": None, "result": "", "branch_ref": None},
        "created_at": int(time.time()),
        "updated_at": int(time.time()),
    }
    tree["nodes"][nid] = node
    _save(_tree_path(args.run_dir), tree)
    kind = "direction" if depth == 1 else "intervention"
    print(f"Added {kind} node {nid} (depth {depth}) under {args.parent}: {args.hypothesis}")


def cmd_baseline(args):
    """Record the evaluated initial artifact once, before any merge decisions."""
    tree = _load_tree(args.run_dir)
    run = _load_run(args.run_dir)
    if run["best_test_score"] is not None or run["best_node"] is not None:
        sys.exit("error: a baseline or best artifact is already recorded; start a new run to change the baseline.")
    root = _node(tree, tree["root"])
    root["metadata"]["test_score"] = args.test_score
    root["metadata"]["branch_ref"] = args.branch_ref
    run["best_node"] = tree["root"]
    run["best_test_score"] = args.test_score
    run["best_branch_ref"] = args.branch_ref
    _stamp(root)
    _save(_tree_path(args.run_dir), tree)
    _save(_run_path(args.run_dir), run)
    print(f"Recorded baseline {tree['root']}: test={args.test_score}, ref={args.branch_ref}")


def cmd_set_status(args):
    tree = _load_tree(args.run_dir)
    node = _node(tree, args.node)
    if args.status not in VALID_STATUS:
        sys.exit(f"error: status must be one of {sorted(VALID_STATUS)}")
    node["status"] = args.status
    _stamp(node)
    _save(_tree_path(args.run_dir), tree)
    print(f"{args.node} -> status={args.status}")


def cmd_set_evidence(args):
    """Write an executor's report back into its node (the Backpropagate step,
    leaf part). Insight propagation upward is a separate, deliberate call."""
    tree = _load_tree(args.run_dir)
    node = _node(tree, args.node)
    if args.status is not None and args.status not in VALID_STATUS:
        sys.exit(f"error: status must be one of {sorted(VALID_STATUS)}")
    meta = node["metadata"]
    if args.dev_score is not None:
        meta["dev_score"] = args.dev_score
    if args.result is not None:
        meta["result"] = args.result
    if args.branch_ref is not None:
        meta["branch_ref"] = args.branch_ref
    if args.insight is not None:
        node["insight"] = args.insight
    node["status"] = args.status or "executed"
    _stamp(node)
    _save(_tree_path(args.run_dir), tree)
    print(f"Wrote evidence to {args.node}: dev_score={meta['dev_score']} status={node['status']}")
    if node["insight"]:
        print(f"  insight: {node['insight']}")
    anc = _ancestors(tree, args.node)
    if anc:
        print(
            "\nReminder: abstract this leaf insight upward. Decide what direction-level "
            f"lesson it implies for ancestors {anc} and record it with "
            f"`tree.py propagate --node {args.node} --insight \"...\"` (and update n0 "
            "global insights if it generalizes)."
        )


def cmd_propagate(args):
    """Backpropagate a distilled, direction-level lesson up the ancestor path.

    The coordinator decides the abstracted wording; this appends it to the
    chosen ancestor(s) so later ideation is conditioned on it. By default it
    updates the immediate parent; --to-root also updates global insights."""
    tree = _load_tree(args.run_dir)
    _node(tree, args.node)  # validate
    targets = _ancestors(tree, args.node)
    if not targets:
        sys.exit("error: node has no ancestors (is it the root?).")
    if not args.to_root:
        targets = targets[:1]  # immediate parent only
    for tid in targets:
        anc = tree["nodes"][tid]
        existing = anc.get("insight", "")
        line = f"[from {args.node}] {args.insight}"
        anc["insight"] = (existing + "\n" + line).strip() if existing else line
        _stamp(anc)
    _save(_tree_path(args.run_dir), tree)
    print(f"Propagated insight from {args.node} to ancestors {targets}")


def cmd_prune(args):
    """Mark a node (and its subtree) pruned. Pruned hypotheses become negative
    constraints -- record *why* so future ideation avoids the dead end."""
    tree = _load_tree(args.run_dir)
    _node(tree, args.node)
    stack = [args.node]
    pruned = []
    while stack:
        cur = stack.pop()
        node = tree["nodes"][cur]
        if node["status"] in ("merged", "root"):
            continue
        node["status"] = "pruned"
        if cur == args.node and args.reason:
            node["metadata"]["prune_reason"] = args.reason
        _stamp(node)
        pruned.append(cur)
        stack.extend(_children(tree, cur))
    _save(_tree_path(args.run_dir), tree)
    print(f"Pruned {pruned}" + (f" -- reason: {args.reason}" if args.reason else ""))


def cmd_merge(args):
    """Record a held-out merge gate decision. Only call this AFTER evaluating
    the candidate on the TEST evaluator in a fresh worktree. Admitting a
    candidate that only improved dev defeats the purpose of the split."""
    tree = _load_tree(args.run_dir)
    run = _load_run(args.run_dir)
    node = _node(tree, args.node)

    direction = run["metric_direction"]
    prev = run["best_test_score"]
    if prev is None:
        sys.exit("error: evaluate the initial artifact and run `tree.py baseline` before merging candidates.")
    if not math.isfinite(prev):
        sys.exit("error: stored best score is not finite; restore a valid scored run before merging.")
    if args.node == tree["root"]:
        sys.exit("error: merge a candidate node; the root stores the initial baseline.")

    improves = args.test_score > prev if direction == "max" else args.test_score < prev
    node["metadata"]["test_score"] = args.test_score
    if args.branch_ref is not None:
        node["metadata"]["branch_ref"] = args.branch_ref
    if improves:
        node["status"] = "merged"
        run["best_node"] = args.node
        run["best_test_score"] = args.test_score
        run["best_branch_ref"] = args.branch_ref or node["metadata"].get("branch_ref")
        _stamp(node)
        _save(_tree_path(args.run_dir), tree)
        _save(_run_path(args.run_dir), run)
        print(
            f"MERGE GATE PASSED: {args.node} test={args.test_score} "
            f"beats previous best={prev}. M_best is now {args.node} "
            f"(ref: {run['best_branch_ref']})."
        )
    else:
        _stamp(node)
        _save(_tree_path(args.run_dir), tree)
        print(
            f"MERGE GATE REJECTED: {args.node} test={args.test_score} does not beat "
            f"best={prev} (direction={direction}). This is informative, not a failure: "
            "a high-dev / low-test gap means the candidate may be exploiting the dev "
            "signal. Record that lesson and keep M_best unchanged."
        )


def cmd_cycle(args):
    """Increment the cycle counter (call once per Observe->Decide pass)."""
    run = _load_run(args.run_dir)
    run["cycles_used"] += 1
    _save(_run_path(args.run_dir), run)
    left = run["budget_cycles"] - run["cycles_used"]
    print(f"Cycle {run['cycles_used']}/{run['budget_cycles']} ({left} remaining).")
    if left <= 0:
        print("Budget exhausted -- finish the run: do a final merge-gate check and report.")


# ----------------------------------------------------------------------------
# Read-only projections
# ----------------------------------------------------------------------------


def _fmt_score(node, run):
    s = node["metadata"].get("dev_score")
    t = node["metadata"].get("test_score")
    bits = []
    if s is not None:
        bits.append(f"dev={s}")
    if t is not None:
        bits.append(f"test={t}")
    return (" [" + " ".join(bits) + "]") if bits else ""


def cmd_observe(args):
    """The Observe step: a compact projection the coordinator re-grounds on at
    the start of each cycle, so decisions come from the tree rather than from a
    lossy conversation history."""
    tree = _load_tree(args.run_dir)
    run = _load_run(args.run_dir)
    nodes = tree["nodes"]
    root = nodes[tree["root"]]

    print("=" * 72)
    print("OBSERVE -- current research state")
    print("=" * 72)
    print(f"Objective       : {run['objective']}")
    print(f"Metric direction: {run['metric_direction']}")
    print(f"Dev evaluator   : {run['dev_eval']}")
    print(f"Test evaluator  : {run['test_eval']}")
    print(
        f"Budget          : cycle {run['cycles_used']}/{run['budget_cycles']}, "
        f"branching {run['branching']}, max depth {run['max_depth']}"
    )
    print(
        f"Current best    : "
        + (
            f"{run['best_node']} (test={run['best_test_score']}, ref={run['best_branch_ref']})"
            if run["best_node"]
            else "none yet -- M_best is the initial material"
        )
    )

    print("\n-- Global insights (root) --")
    print(root["insight"].strip() if root["insight"].strip() else "  (none yet)")

    # Active frontier = pending/running leaves
    frontier = [
        n
        for n in nodes.values()
        if n["status"] in ("pending", "running") and not _children(tree, n["id"])
    ]
    print("\n-- Active frontier (selectable hypotheses) --")
    if not frontier:
        print("  (empty -- ideate new children under a promising node)")
    for n in sorted(frontier, key=lambda x: _id_sort_key(x["id"])):
        anc = _ancestors(tree, n["id"])
        anc_ins = " | ".join(
            nodes[a]["insight"].replace("\n", " ")[:80] for a in anc if nodes[a]["insight"].strip()
        )
        print(f"  {n['id']} (depth {n['depth']}, {n['status']}): {n['hypothesis']}")
        if anc_ins:
            print(f"      ancestor insights: {anc_ins}")

    # Validated / executed leaves with evidence
    executed = [n for n in nodes.values() if n["status"] in ("executed", "merged")]
    print("\n-- Executed / merged nodes (evidence) --")
    if not executed:
        print("  (none yet)")
    for n in sorted(executed, key=lambda x: _id_sort_key(x["id"])):
        print(f"  {n['id']} [{n['status']}]{_fmt_score(n, run)}: {n['hypothesis']}")
        if n["insight"].strip():
            print(f"      insight: {n['insight'].splitlines()[0][:120]}")

    # Pruned lessons (negative constraints)
    pruned = [n for n in nodes.values() if n["status"] == "pruned"]
    print("\n-- Pruned lessons (negative constraints -- avoid these) --")
    if not pruned:
        print("  (none yet)")
    for n in sorted(pruned, key=lambda x: _id_sort_key(x["id"])):
        reason = n["metadata"].get("prune_reason", "")
        print(f"  {n['id']}: {n['hypothesis']}" + (f" -- {reason}" if reason else ""))

    print("\n" + "=" * 72)
    print(
        "Next: Ideate children under a promising node, Select one or more frontier\n"
        "leaves, Dispatch each to an executor subagent in an isolated worktree."
    )


def cmd_status(args):
    """ASCII tree render -- useful for reports and quick scans."""
    tree = _load_tree(args.run_dir)
    run = _load_run(args.run_dir)
    nodes = tree["nodes"]

    symbol = {
        "root": "*",
        "pending": "o",
        "running": "~",
        "executed": "=",
        "merged": "V",
        "pruned": "x",
    }

    def render(nid, prefix=""):
        n = nodes[nid]
        sym = symbol.get(n["status"], "?")
        best = " <== M_best" if nid == run["best_node"] else ""
        print(f"{prefix}[{sym}] {nid} {n['hypothesis'][:70]}{_fmt_score(n, run)}{best}")
        kids = sorted(_children(tree, nid), key=_id_sort_key)
        for i, c in enumerate(kids):
            render(c, prefix + "    ")

    print(f"Run: {run['objective']}")
    print(f"Cycle {run['cycles_used']}/{run['budget_cycles']}  |  legend: * root  o pending  ~ running  = executed  V merged  x pruned\n")
    render(tree["root"])


def cmd_validate(args):
    """Check invariants -- catch a corrupted or inconsistent tree early."""
    tree = _load_tree(args.run_dir)
    run = _load_run(args.run_dir)
    nodes = tree["nodes"]
    problems = []
    if tree["root"] not in nodes:
        problems.append("root id not present in nodes")
    for nid, n in nodes.items():
        p = n.get("parent")
        if p is not None and p not in nodes:
            problems.append(f"{nid}: parent {p} missing")
        if n["status"] not in VALID_STATUS:
            problems.append(f"{nid}: invalid status {n['status']}")
    if run["best_node"] and run["best_node"] not in nodes:
        problems.append(f"best_node {run['best_node']} missing")
    merged = [nid for nid, n in nodes.items() if n["status"] == "merged"]
    if run["best_node"] and run["best_node"] not in merged and run["best_node"] != tree["root"]:
        problems.append(f"best_node {run['best_node']} is not marked merged")
    if problems:
        print("INVALID:")
        for p in problems:
            print(f"  - {p}")
        sys.exit(1)
    print(f"OK -- {len(nodes)} nodes, root={tree['root']}, best={run['best_node']}")


# ----------------------------------------------------------------------------
# CLI
# ----------------------------------------------------------------------------


def build_parser():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--run-dir", default=".", help="Run directory holding .arbor/ (default: current dir)")
    sub = p.add_subparsers(dest="command", required=True)

    s = sub.add_parser("init", help="Initialize a new AO run / hypothesis tree")
    s.add_argument("--objective", required=True, help="Natural-language research objective (root hypothesis)")
    s.add_argument("--dev-eval", required=True, help="Command/description of the development evaluator")
    s.add_argument("--test-eval", required=True, help="Command/description of the held-out test evaluator")
    s.add_argument("--material", default="", help="Path/ref to the initial artifact M_0")
    s.add_argument("--metric-direction", choices=["max", "min"], default="max", help="Is higher or lower the better score?")
    s.add_argument("--branching", type=int, default=3, help="Max children proposed per parent (k)")
    s.add_argument("--max-depth", type=int, default=2, help="Max tree depth (directions at 1, interventions at 2+)")
    s.add_argument("--budget", type=int, default=20, help="Coordinator cycle budget B")
    s.add_argument("--force", action="store_true", help="Overwrite an existing run")
    s.set_defaults(func=cmd_init)

    s = sub.add_parser("baseline", help="Record the initial artifact's held-out score before candidate merges")
    s.add_argument("--test-score", type=_finite_score, required=True)
    s.add_argument("--branch-ref", required=True, help="Immutable git commit or named branch of M_0")
    s.set_defaults(func=cmd_baseline)

    s = sub.add_parser("observe", help="Print the research-state projection (start of each cycle)")
    s.set_defaults(func=cmd_observe)

    s = sub.add_parser("add-node", help="Add a pending child hypothesis under a parent (Ideate)")
    s.add_argument("--parent", required=True, help="Parent node id (n0 for a new research direction)")
    s.add_argument("--hypothesis", required=True, help="Falsifiable claim this node tests")
    s.set_defaults(func=cmd_add_node)

    s = sub.add_parser("set-status", help="Set a node's status manually")
    s.add_argument("--node", required=True)
    s.add_argument("--status", required=True, help=f"One of {sorted(VALID_STATUS)}")
    s.set_defaults(func=cmd_set_status)

    s = sub.add_parser("set-evidence", help="Write an executor report into its node (Backpropagate, leaf)")
    s.add_argument("--node", required=True)
    s.add_argument("--dev-score", type=_finite_score, default=None, help="Dev evaluator score returned by the executor")
    s.add_argument("--result", default=None, help="Factual result summary")
    s.add_argument("--insight", default=None, help="Distilled, reusable lesson from this experiment")
    s.add_argument("--branch-ref", default=None, help="Git branch/commit/worktree path of the artifact")
    s.add_argument("--status", default=None, help=f"Override status (default: executed), one of {sorted(VALID_STATUS)}")
    s.set_defaults(func=cmd_set_evidence)

    s = sub.add_parser("propagate", help="Abstract a leaf insight up to ancestors (Backpropagate, upward)")
    s.add_argument("--node", required=True, help="The leaf the lesson came from")
    s.add_argument("--insight", required=True, help="Direction-level abstraction of the lesson")
    s.add_argument("--to-root", action="store_true", help="Also record as a global insight on the root")
    s.set_defaults(func=cmd_propagate)

    s = sub.add_parser("prune", help="Prune a falsified node and its subtree (Decide)")
    s.add_argument("--node", required=True)
    s.add_argument("--reason", default="", help="Why this direction is a dead end (becomes a negative constraint)")
    s.set_defaults(func=cmd_prune)

    s = sub.add_parser("merge", help="Record a held-out merge gate decision (Decide)")
    s.add_argument("--node", required=True)
    s.add_argument("--test-score", type=_finite_score, required=True, help="Score on the TEST evaluator in a fresh worktree")
    s.add_argument("--branch-ref", default=None, help="Artifact ref to promote if it passes")
    s.set_defaults(func=cmd_merge)

    s = sub.add_parser("cycle", help="Increment the coordinator cycle counter")
    s.set_defaults(func=cmd_cycle)

    s = sub.add_parser("status", help="Render the tree as ASCII (for reports)")
    s.set_defaults(func=cmd_status)

    s = sub.add_parser("validate", help="Check tree invariants")
    s.set_defaults(func=cmd_validate)

    return p


def main():
    args = build_parser().parse_args()
    args.func(args)


if __name__ == "__main__":
    main()