"""argparse and the commands; each launch is record -> one `claude -p` call -> record."""
import argparse
import json
import os
import re
import string
import subprocess
import sys
import time
import uuid
from pathlib import Path

from grm import call, launch, store, stream

SKILL_DIR = Path(__file__).resolve().parent.parent
NAME_RE = re.compile(r"^[A-Za-z0-9._-]+$")
NOTE_SLACK_S = 2.0     # coarse file times: a note written right at the start still counts
EXCERPT_CHARS = 500
LAUNCHES = ("train", "spawn", "rebase")


class Refusal(Exception):
    """Refused before or instead of a launch: exit 2, payload printed as the JSON object."""

    def __init__(self, error: str, **fields):
        super().__init__(error)
        self.payload = {"error": error, **fields}


class _Parser(argparse.ArgumentParser):
    def error(self, message):
        raise Refusal("usage", detail=message)


def _emit(obj: dict) -> None:
    sys.stdout.write(json.dumps(obj) + "\n")
    sys.stdout.flush()


def render_template(name: str, values: dict) -> str:
    with open(SKILL_DIR / "templates" / name, "r", encoding="utf-8") as f:
        return string.Template(f.read()).safe_substitute(values)


def git(args, cwd=None, timeout=60) -> subprocess.CompletedProcess:
    """Every helper git call: `git --no-optional-locks …`, output captured."""
    return subprocess.run(["git", "--no-optional-locks", *args], cwd=cwd, stdin=subprocess.DEVNULL,
                          capture_output=True, timeout=timeout)


def _git_out(args, cwd) -> str | None:
    try:
        r = git(args, cwd=cwd)
    except (OSError, subprocess.SubprocessError):
        return None
    out = r.stdout.decode("utf-8", "replace").strip()
    return out if r.returncode == 0 and out else None


def git_head(cwd) -> str | None:
    return _git_out(["rev-parse", "HEAD"], cwd)


def delta_header(cwd, base_commit: str | None) -> str:
    """Commits since the base's commit plus the working tree's shortstat against it."""
    if not base_commit:
        return ("No base commit is recorded (the base's cwd was not a git repository, or git failed "
                "at training). Treat everything on disk as possibly changed since training.")
    parts = []
    try:
        lg = git(["log", "--oneline", "%s..HEAD" % base_commit], cwd=cwd)
        st = git(["diff", "--shortstat", base_commit], cwd=cwd)
        log_txt = lg.stdout.decode("utf-8", "replace").strip() if lg.returncode == 0 else \
            "(git log failed: %s)" % lg.stderr.decode("utf-8", "replace").strip()
        st_txt = st.stdout.decode("utf-8", "replace").strip() if st.returncode == 0 else \
            "(git diff failed: %s)" % st.stderr.decode("utf-8", "replace").strip()
    except (OSError, subprocess.SubprocessError) as e:
        log_txt = st_txt = "(git failed: %s)" % e
    parts.append("Base commit: `%s`." % base_commit)
    parts.append("Commits since then (`git log --oneline %s..HEAD`):\n\n```\n%s\n```"
                 % (base_commit[:12], log_txt or "(none)"))
    parts.append("Working tree against the base commit (`git diff --shortstat %s`):\n\n```\n%s\n```"
                 % (base_commit[:12], st_txt or "(no changes)"))
    return "\n\n".join(parts)


def _name(kind: str, value: str) -> str:
    if not NAME_RE.match(value or ""):
        raise Refusal("bad %s name" % kind, **{kind: value})
    return value


def _read_brief(path: str) -> str:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read()
    except OSError as e:
        raise Refusal("cannot read brief", brief=path, detail=str(e))


def _load_run(run: str) -> tuple[Path, dict]:
    """State dir and run.json of a launch; the stored binary must still exist (no job file yet)."""
    state = store.state_dir(run)
    r = store.read_json(state / "run.json")
    if r is None:
        raise Refusal("no run.json for this run; call init first", state=str(state))
    if not os.path.isfile(r["argv"][0]):
        raise Refusal("claude binary missing", binary=r["argv"][0])
    return state, r


def _load_base(state: Path, b: str) -> dict:
    base = store.read_json(state / "bases" / (b + ".json"))
    if base is None:
        raise Refusal("unknown base", base=b, bases=sorted(p.stem for p in (state / "bases").glob("*.json")))
    return base


def _warmth(lines: list, base: dict, now: float) -> dict:
    """B6: from the ledger line of this base whose call read the base prefix last."""
    hits = [ln for ln in lines if ln.get("base") == base["base"] and stream.parse_iso(ln.get("hit_at"))]
    if not hits:
        return {**stream.warmth(None, None, now), "ttl_s": base.get("ttl_s")}
    last = max(hits, key=lambda ln: stream.parse_iso(ln["hit_at"]))
    ttl = stream.infer_ttl(last.get("tokens"), base.get("ttl_s")) or stream.TTL_1H
    return {**stream.warmth(stream.parse_iso(last["hit_at"]), ttl, now), "ttl_s": ttl}


def _execute(run: str, state: Path, r: dict, *, job, kind, base, sid, argv, prompt,
             before=(None, 0.0), base_ctx=None):
    """One call: job file `running` -> run_claude -> job file final + one ledger line.
    Returns (job record, StreamState|None, usage|None, ledger line|None)."""
    fp = r["fingerprint"]
    jpath = state / "jobs" / (job + ".json")
    raw = state / "raw" / (job + ".jsonl")
    store.atomic_write_bytes(state / "briefs" / (job + ".md"), prompt.encode("utf-8"))
    t0 = time.time()
    rec = {"job": job, "kind": kind, "base": base, "session_id": sid, "status": "running",
           "reason": None, "started": stream.iso(t0), "ended": None, "child_pid": None,
           "raw": str(raw), "note": None, "cache": None, "result_excerpt": None}
    store.write_json(jpath, rec)

    def on_start(pid):
        rec["child_pid"] = pid
        store.write_json(jpath, rec)
    try:
        ss, code, stopped = call.run_claude(argv, prompt, launch.scrubbed_env(fp["binary"]),
                                            fp["cwd"], raw, on_start)
    except OSError as e:
        rec.update(status="error", reason="launch_error", ended=stream.iso(),
                   result_excerpt=str(e)[:EXCERPT_CHARS])
        store.write_json(jpath, rec)
        return rec, None, None, None
    wall = time.time() - t0
    status, reason = stream.classify(ss, code, stopped)
    u = stream.usage_from_result(ss.result, *before) if ss.result else None
    cache = stream.cache_verdict(ss.first_usage, base_ctx) if kind != "train" else None
    hit_t = None
    if kind in ("train", "rebase") and status == "ok":
        hit_t = ss.last_assistant_at
    elif kind == "spawn" and cache in ("hit", "partial"):
        hit_t = ss.first_at
    note = None
    if kind == "spawn":
        p = state / "notes" / (job + ".md")
        if p.is_file() and p.stat().st_mtime >= t0 - NOTE_SLACK_S:
            note = str(p)
    text = (ss.result or {}).get("result")
    rec.update(session_id=ss.session_id or sid, status=status, reason=reason, ended=stream.iso(),
               cache=cache, note=note,
               result_excerpt=text[:EXCERPT_CHARS] if isinstance(text, str) else None)
    store.write_json(jpath, rec)
    models = u["models_used"] if u else []
    line = {"ts": rec["ended"], "run": r["run"], "job": job, "kind": kind, "base": base,
            "session_id": rec["session_id"], "model": models[0] if models else None,
            "models_used": models,
            "tokens": u["tokens"] if u else {k: None for k in stream.TOKEN_KEYS},
            "cost_delta_usd": u["cost_delta"] if u else None,
            "cost_cumulative_usd": u["total"] if u else None, "cache": cache, "status": status,
            "wall_s": round(wall, 1), "hit_at": stream.iso(hit_t) if hit_t else None}
    store.ledger_append(run, line)
    return rec, ss, u, line


def _cost_basis(ss) -> float | None:
    """The base's cost basis stays unrounded: deltas against a rounded basis drift by 1e-6."""
    total = (ss.result or {}).get("total_cost_usd")
    return None if total is None else float(total)


def _out(rec: dict, line: dict | None) -> dict:
    out = dict(rec)
    if line:
        out.update(model=line["model"], models_used=line["models_used"], tokens=line["tokens"],
                   cost=line["cost_delta_usd"], cost_cumulative_usd=line["cost_cumulative_usd"])
    return out


def cmd_init(args) -> dict:
    """B1: validate, then write the frozen preamble copy, the empty MCP config and run.json once."""
    state = store.state_dir(args.run)
    if (state / "run.json").exists():
        raise Refusal("run.json exists; init never overwrites a run", state=str(state))
    binary = launch.resolve_binary(args.binary)
    if not binary or not os.path.isfile(binary):
        raise Refusal("no claude binary (no --binary, no $CLAUDE_CODE_EXECPATH, none on PATH)",
                      binary=binary)
    cwd = os.path.abspath(args.cwd) if args.cwd else _git_out(["rev-parse", "--show-toplevel"], os.getcwd())
    if not cwd or not os.path.isdir(cwd):
        raise Refusal("no usable cwd (pass --cwd, or run inside a git repository)", cwd=cwd)
    if not launch.MODEL_RE.match(args.model):     # before anything is written
        raise Refusal("model must be a full id (aliases are refused)", model=args.model)
    state = store.init_state(args.run)
    pre = state / "preamble.md"
    store.atomic_write_bytes(pre, (SKILL_DIR / "preamble.md").read_bytes())
    mcp = state / "mcp-empty.json"
    store.write_json(mcp, {"mcpServers": {}})
    fpd = launch.build_fingerprint(binary=binary, model=args.model, effort=args.effort,
                                   cwd=os.path.abspath(cwd), agent=args.agent,
                                   preamble_path=str(pre), mcp_config=str(mcp))
    run = {"run": os.path.abspath(args.run), "created": stream.iso(), **fpd}
    store.write_json(state / "run.json", run)
    return {"state": str(state), **run}


def cmd_train(args) -> dict:
    """B1: one blocking call with `--session-id`; the base file only on `ok`."""
    b = _name("base", args.base)
    state, r = _load_run(args.run)
    bpath = state / "bases" / (b + ".json")
    if bpath.exists():
        raise Refusal("base exists", base=b)
    brief = _read_brief(args.brief)
    sid, job = str(uuid.uuid4()), b + "-train"
    prompt = render_template("train.md", {
        "base": b, "run": r["run"], "state_dir": str(state), "notes_dir": str(state / "notes"),
        "logs_dir": str(state / "logs"), "brief": brief})
    rec, ss, u, line = _execute(args.run, state, r, job=job, kind="train", base=b, sid=sid,
                                argv=launch.train_argv(r, sid), prompt=prompt)
    if rec["status"] == "ok":
        store.write_json(bpath, {
            "base": b, "session_id": rec["session_id"], "train_job": job,
            "commit": git_head(r["fingerprint"]["cwd"]), "trained_at": rec["ended"],
            "context_tokens": stream.context_tokens(ss.last_usage), "model_usage": u["model_usage"],
            "cost_cumulative_usd": _cost_basis(ss), "ttl_s": stream.infer_ttl(u["tokens"], None),
            "rebases": []})
    return _out(rec, line)


def cmd_spawn(args) -> dict:
    """B2, B3, B5: one blocking fork of the base; never writes the base file."""
    b, j = _name("base", args.base), _name("job", args.job)
    state, r = _load_run(args.run)
    base = _load_base(state, b)
    if (state / "jobs" / (j + ".json")).exists():
        raise Refusal("job exists; pick a new --job", job=j)
    brief = _read_brief(args.brief)
    warm = _warmth(store.ledger_lines(args.run), base, time.time())
    sid = str(uuid.uuid4())
    prompt = render_template("spawn.md", {
        "block": j, "base": b, "run": r["run"],
        "delta": delta_header(r["fingerprint"]["cwd"], base.get("commit")),
        "note_path": str(state / "notes" / (j + ".md")), "logs_dir": str(state / "logs"),
        "brief": brief})
    rec, _, _, line = _execute(args.run, state, r, job=j, kind="spawn", base=b, sid=sid,
                               argv=launch.fork_argv(r, base["session_id"], sid), prompt=prompt,
                               before=(base.get("model_usage"), base.get("cost_cumulative_usd")),
                               base_ctx=base.get("context_tokens"))
    return {**_out(rec, line), "base_state": warm["base_state"]}


def cmd_rebase(args) -> dict:
    """B7: resume the base without a fork; on `ok` the base takes the result's session id."""
    b = _name("base", args.base)
    state, r = _load_run(args.run)
    base = _load_base(state, b)
    brief = _read_brief(args.brief)
    running = sorted(rec["job"] for rec in map(store.read_json, (state / "jobs").glob("*.json"))
                     if rec and rec.get("base") == b and rec.get("status") == "running")
    warnings = ["jobs of %s still marked running: %s" % (b, ", ".join(running))] if running else []
    job = "%s-rebase%d" % (b, len(base.get("rebases") or []) + 1)
    cwd = r["fingerprint"]["cwd"]
    prompt = render_template("rebase.md", {
        "base": b, "run": r["run"], "notes_dir": str(state / "notes"),
        "delta": delta_header(cwd, base.get("commit")), "brief": brief})
    rec, ss, u, line = _execute(args.run, state, r, job=job, kind="rebase", base=b, sid=None,
                                argv=launch.resume_argv(r, base["session_id"]), prompt=prompt,
                                before=(base.get("model_usage"), base.get("cost_cumulative_usd")),
                                base_ctx=base.get("context_tokens"))
    if rec["status"] == "ok":
        base["rebases"] = list(base.get("rebases") or []) + [{
            "job": job, "at": rec["ended"], "from_commit": base.get("commit"),
            "from_session": base["session_id"]}]
        base.update(session_id=rec["session_id"] or base["session_id"], commit=git_head(cwd),
                    context_tokens=stream.context_tokens(ss.last_usage), model_usage=u["model_usage"],
                    cost_cumulative_usd=_cost_basis(ss), ttl_s=stream.infer_ttl(u["tokens"], base.get("ttl_s")))
        store.write_json(state / "bases" / (b + ".json"), base)
    return {**_out(rec, line), "warnings": warnings}


def cmd_status(args) -> dict:
    """B6 (warmth per base) and B3 (re-reading a job's result)."""
    state = store.state_dir(args.run)
    if not (state / "run.json").exists():
        raise Refusal("no run.json for this run; call init first", state=str(state))
    if args.job:
        rec = store.read_json(state / "jobs" / (_name("job", args.job) + ".json"))
        if rec is None:
            raise Refusal("unknown job", job=args.job)
        return rec
    lines, now = store.ledger_lines(args.run), time.time()
    bases = {}
    for p in sorted((state / "bases").glob("*.json")):
        base = store.read_json(p)
        bases[base["base"]] = {**{k: base.get(k) for k in ("session_id", "commit", "context_tokens")},
                               **_warmth(lines, base, now)}
    jobs = [store.read_json(p) for p in sorted((state / "jobs").glob("*.json"))]
    return {"state": str(state), "bases": bases, "jobs": jobs}


COMMANDS = {"init": cmd_init, "train": cmd_train, "spawn": cmd_spawn, "rebase": cmd_rebase,
            "status": cmd_status}


def _parser() -> argparse.ArgumentParser:
    p = _Parser(prog="gremlins.py", description="spire:gremlins helper; prints one JSON object.")
    sub = p.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("init")
    s.add_argument("--model", required=True)
    s.add_argument("--effort", default="high", choices=launch.EFFORTS)
    s.add_argument("--cwd")
    s.add_argument("--agent")
    s.add_argument("--binary")
    for name in LAUNCHES:
        s2 = sub.add_parser(name)
        s2.add_argument("--base", required=True)
        if name == "spawn":
            s2.add_argument("--job", required=True)
        s2.add_argument("--brief", required=True)
    sst = sub.add_parser("status")
    sst.add_argument("--job")
    for sp in sub.choices.values():
        sp.add_argument("--run", required=True)
    return p


def main(argv=None) -> int:
    """Exit 0 ok, 1 a call that ended `error`/`stopped`, 2 refused or a usage error."""
    try:
        args = _parser().parse_args(argv)
        out = COMMANDS[args.cmd](args)
    except Refusal as e:
        _emit(e.payload)
        return 2
    except (ValueError, OSError) as e:
        _emit({"error": str(e), "type": type(e).__name__})
        return 2
    _emit(out)
    return 1 if args.cmd in LAUNCHES and out.get("status") != "ok" else 0