#!/usr/bin/env python3
"""Carry a cut_spec's file_changes anchors from the spec's base to a worktree.

Usage: context-pack.py <cut_spec id> <worktree> [--since <sha>] [--only <k>|r<k>[,...]]
                       [--spec-json <file>] [--max-lines <n>]

Prints a header, a status table (one line per anchor: file_changes as <k>, reads
as r<k>) and bounded windows of the code each anchor names, numbered in the
worktree's lines: edit windows, then READ windows, then one caller list per
distinct symbol. A symbol an anchor names must be in its mapped window or the
status gains SYMBOL_MISSING. Exit: 0 every anchor carried; 3 an anchor cannot be
carried (OUT_OF_RANGE, UNRESOLVED, MISSING read, SYMBOL_MISSING, GONE outside
the spec's deletes); 2 usage; 4 spec unreadable. Standard library plus git.
Errors name a field and a code, never an input value.
--max-lines bounds the windows and caller lists, not the table: the table always
lists every anchor, so a budget smaller than the table prints the table and no
window. Precedence: table, edit windows, read windows, caller lists.
"""
import argparse
import json
import os
import re
import shlex
import subprocess
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mind  # noqa: E402

WINDOW_CAP = 100
CONTEXT = 3
DEFAULT_MAX_LINES = 800
HUNK = re.compile(r"^@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@")
MAX_CALLER_SYMBOLS = 8
MAX_CALLER_HITS = 50
MAX_CALLER_FILES = 20
HIT_TEXT = 160
ONLY = re.compile(r"^(r?)([0-9]+)$")
IDENT = re.compile(r"^[A-Za-z0-9_]+$")
HEX = re.compile(r"^[0-9a-fA-F]{7,40}$")


class Fail(Exception):
    def __init__(self, code, msg):
        super().__init__(msg)
        self.code, self.msg = code, msg


class Parser(argparse.ArgumentParser):
    def error(self, message):  # argparse would echo the offending value
        sys.stderr.write("context-pack: usage error\n" + self.format_usage())
        sys.exit(2)


def git(wt, *args):
    """The one git call: stdout bytes, or None when git failed; stderr is always
    discarded. Paths are literal pathspecs here, so a spec path holding [ ] * ?
    names that file and no sibling."""
    try:
        p = subprocess.run(["git", "-C", wt, "--literal-pathspecs", *args], stdout=subprocess.PIPE,
                           stderr=subprocess.DEVNULL, stdin=subprocess.DEVNULL)
    except OSError:
        return None
    return p.stdout if p.returncode == 0 else None


def text_lines(raw):
    s = raw.decode("utf-8", "replace").replace("\r\n", "\n")
    parts = s.split("\n")
    if parts[-1] == "":
        parts.pop()
    return parts


# ---- the spec ------------------------------------------------------------

def spec_from_json(obj):
    if isinstance(obj, dict) and "document" in obj:
        obj = obj["document"]
    if isinstance(obj, dict) and "kind" in obj and "value" in obj:
        obj = obj["value"]
    if not isinstance(obj, dict) or "base" not in obj or "file_changes" not in obj:
        raise Fail(4, "spec: not a cut_spec value")
    return check_spec_shape(obj)


def is_int(v):
    return isinstance(v, int) and not isinstance(v, bool)


def check_spec_shape(spec):
    """A wrong-typed field is a typed refusal naming the field and InvalidInput,
    never the value."""
    def bad(field):
        raise Fail(4, "spec: %s: InvalidInput" % field)
    if not isinstance(spec["base"], str):
        bad("base")
    def check_loc(at, loc):
        if not isinstance(loc, dict):
            bad(at)
        if not isinstance(loc.get("path"), str):
            bad(at + ".path")
        if not is_int(loc.get("line")):
            bad(at + ".line")
        if loc.get("end_line") is not None and not is_int(loc["end_line"]):
            bad(at + ".end_line")
        if loc.get("symbol") is not None and not isinstance(loc["symbol"], str):
            bad(at + ".symbol")
    fcs = spec["file_changes"]
    if not isinstance(fcs, list):
        bad("file_changes")
    for k, fc in enumerate(fcs, 1):
        at = "file_changes[%d]" % k
        if not isinstance(fc, dict):
            bad(at)
        check_loc(at + ".location", fc.get("location"))
        if "change" in fc and not isinstance(fc["change"], str):
            bad(at + ".change")
    reads = spec.get("reads")
    if reads is None:
        reads = []
    if not isinstance(reads, list):
        bad("reads")
    for k, rd in enumerate(reads, 1):
        at = "reads[%d]" % k
        if not isinstance(rd, dict):
            bad(at)
        check_loc(at + ".location", rd.get("location"))
        if "why" in rd and not isinstance(rd["why"], str):
            bad(at + ".why")
    dels = spec.get("deletes", [])
    if not isinstance(dels, list):
        bad("deletes")
    for k, d in enumerate(dels, 1):
        if not isinstance(d, dict) or not isinstance(d.get("path"), str):
            bad("deletes[%d].path" % k)
    return spec


def mcp_view(spec_id):
    if ":cut_spec:" not in spec_id:
        raise Fail(2, "spec id: not a cut_spec id")
    try:
        view = mind.retry_once(lambda: mind.call("view", {"id": {"kind": "CutSpec", "id": spec_id}}, 60))
    except mind.MindError as e:
        raise Fail(4, ("view" if e.kind in ("tool", "answer") else e.kind) + ": " + e.detail)
    if view is None:
        raise Fail(4, "view: no such cut_spec")
    return spec_from_json(view)


def load_spec(spec_id, spec_json):
    try:
        if spec_json:
            with open(spec_json, encoding="utf-8") as f:
                return spec_from_json(json.load(f))
        return mcp_view(spec_id)
    except Fail as e:
        if e.code == 4:
            raise Fail(4, e.msg + "; --spec-json reads a saved view")
        raise
    except (OSError, ValueError):
        raise Fail(4, "--spec-json: unreadable; --spec-json reads a saved view")


# ---- carrying an anchor --------------------------------------------------

def carry(s, e, hunks):
    """Map old lines s..e through -U0 hunks (a,b,c,d). Returns (ns, ne, ins, dele, touched)."""
    shift = net_in = ins = dele = 0
    touched = False
    lo = hi = None
    for a, b, c, d in hunks:
        if b == 0:                      # insertion after old line a
            if a < s:
                shift += d
                continue
            if a >= e:                  # at the last line or after: outside
                continue
            touched = True
            ins += d
            net_in += d
        else:
            last = a + b - 1
            if last < s:
                shift += d - b
                continue
            if a > e:
                continue
            touched = True
            ins += d
            dele += b
            net_in += d - b
            if a < s:                   # the first surviving line: c, or c+1 after a pure deletion
                lo = c if d > 0 else c + 1
            if last > e:                # the last surviving line: c+d-1, or c after a pure deletion
                hi = c + d - 1 if d > 0 else c
    ns, ne = s + shift, e + shift + net_in
    if lo is not None:
        ns = min(ns, lo)
    if hi is not None:
        ne = max(ne, hi)
    return ns, max(ne, ns), ins, dele, touched


class Tree:
    def __init__(self, wt, base):
        self.wt, self.base = wt, base
        self.changes = {}
        self._hunks, self._basefile, self._wtfile = {}, {}, {}

    def load_changes(self):
        out = git(self.wt, "diff", "-M", "--name-status", "-z", "--no-color", "--no-ext-diff", self.base)
        if out is None:
            return False
        parts = out.decode("utf-8", "replace").split("\0")
        i = 0
        while i < len(parts) and parts[i]:
            code = parts[i]
            if code[0] in "RC" and i + 2 < len(parts):
                self.changes[parts[i + 1]] = (code[0], parts[i + 2])
                i += 3
            else:
                self.changes[parts[i + 1]] = (code[0], None)
                i += 2
        return True

    def hunks(self, since, *paths):
        key = (since, paths)
        if key not in self._hunks:
            out = git(self.wt, "diff", "-M", "-U0", "--no-color", "--no-ext-diff", since, "--", *paths)
            hs = None
            if out is not None:
                hs = []
                for ln in out.decode("utf-8", "replace").split("\n"):
                    m = HUNK.match(ln)
                    if m:
                        a, b, c, d = m.groups()
                        hs.append((int(a), 1 if b is None else int(b), int(c), 1 if d is None else int(d)))
            self._hunks[key] = hs
        return self._hunks[key]

    def base_lines(self, path):
        """None when the path is absent at base."""
        if path not in self._basefile:
            raw = git(self.wt, "show", self.base + ":" + path)
            self._basefile[path] = None if raw is None else text_lines(raw)
        return self._basefile[path]

    def wt_lines(self, path):
        if path not in self._wtfile:
            try:
                with open(os.path.join(self.wt, *path.split("/")), "rb") as f:
                    self._wtfile[path] = text_lines(f.read())
            except OSError:
                self._wtfile[path] = None
        return self._wtfile[path]


def safe_path(path):
    return bool(path) and not path.startswith(("/", "\\")) and ":" not in path \
        and ".." not in path.replace("\\", "/").split("/")


def has_symbol(symbol, lines):
    text = "\n".join(lines)
    if IDENT.match(symbol):
        return re.search(r"(?<![A-Za-z0-9_])" + re.escape(symbol) + r"(?![A-Za-z0-9_])", text) is not None
    return symbol in text


def resolve(tree, label, fc, deleted_ok, base_ok, listed, read=False):
    """One anchor: carried to the worktree, then its symbol (if any) looked up in
    the mapped window (the AT BASE window for GONE)."""
    a = carry_anchor(tree, label, fc, deleted_ok, base_ok, listed, read)
    sym = fc["location"].get("symbol")
    a["symbol"] = sym or None
    if sym and a["window"]:
        kind, path, s, e = a["window"]
        lines = tree.base_lines(path) if kind == "base" else tree.wt_lines(path)
        if lines is not None and not has_symbol(sym, lines[s - 1:e]):
            a["status"] += " SYMBOL_MISSING " + sym
            a["exit"] = 3
    return a


def carry_anchor(tree, label, fc, deleted_ok, base_ok, listed, read):
    loc = fc["location"]
    path, s = loc["path"], int(loc["line"])
    e = int(loc.get("end_line") or s)
    a = {"k": label, "path": path, "s": s, "e": e, "change": str(fc.get("change", "")),
         "rank": 1, "exit": 0, "status": "", "now": None, "window": None, "read": read}
    if not base_ok:
        a.update(rank=0, exit=3, status="UNRESOLVED base not a commit here: git fetch origin " + tree.base)
        return a
    if not listed or not safe_path(path):
        a.update(rank=0, exit=3, status="UNRESOLVED path or git diff unusable")
        return a
    base = tree.base_lines(path)
    if base is None and read:
        a.update(rank=0, exit=3, status="MISSING (read path absent at base)")
        return a
    if base is None:
        sib = [p for p in tree.siblings() if p.rsplit("/", 1)[-1] == path.rsplit("/", 1)[-1]][:3]
        a.update(rank=0, status="NEW" + (" (at base: " + ", ".join(sib) + ")" if sib else ""))
        a["now"] = (path, s, e)
        if tree.wt_lines(path) is not None:
            a["window"] = ("wt", path, s, e)
        return a
    code, newpath = tree.changes.get(path, ("", None))
    if code == "D":
        ok = path in deleted_ok
        a.update(rank=0, exit=0 if ok else 3, status="GONE" + ("" if ok else " (not named by the spec's deletes)"))
        a["window"] = ("base", path, s, e)
        return a
    if s > len(base) or e > len(base):
        a.update(rank=0, exit=3, status="OUT_OF_RANGE (base file has %d lines)" % len(base))
        return a
    now = newpath or path
    hs = tree.hunks(tree.base, *dict.fromkeys([path, now]))
    if hs is None:
        a.update(rank=0, exit=3, status="UNRESOLVED git diff failed")
        return a
    ns, ne, ins, dele, touched = carry(s, e, hs)
    if touched:
        st = "TOUCHED(+%d/-%d inside)" % (ins, dele)
    elif ns != s:
        st = "MOVED(%+d)" % (ns - s)
    else:
        st = "SAME"
    if newpath:
        st = "RENAMED -> " + newpath + " " + st
    a.update(status=st, now=(now, ns, ne))
    a["window"] = ("wt", now, ns, ne)
    return a


def table_line(a):
    line = "[%s] %s:%d-%d  %s" % (a["k"], a["path"], a["s"], a["e"], a["status"])
    if a["now"]:
        line += "  -> %s:%d-%d" % a["now"]
    return line + "  # " + " ".join(a["change"].split())[:70]


def window_body(tree, a):
    kind, path, s, e = a["window"]
    lines = tree.base_lines(path) if kind == "base" else tree.wt_lines(path)
    where = "%s:%d-%d%s  %s" % (path, s, e, " AT BASE" if kind == "base" else "", a["status"])
    cap = ("== READ %s: %s  (%s)" % (a["k"], a["change"][:200], where)) if a["read"] else "== [%s] %s" % (a["k"], where)
    if lines is None:
        return cap, ["(file not readable in the worktree)"]
    lo, hi = max(1, s - CONTEXT), min(len(lines), e + CONTEXT)
    body = ["%6d  %s" % (n, lines[n - 1]) for n in range(lo, hi + 1)]
    return cap, body or ["(no lines: file has %d lines)" % len(lines)]


def overlaps_since(tree, since, a):
    kind, path, s, e = a["window"]
    if kind != "wt":
        return False
    hs = tree.hunks(since, path)
    if hs is None:
        return False
    for _, _, c, d in hs:
        if d == 0:
            if s - 1 <= c <= e:
                return True
        elif c <= e and c + d - 1 >= s:
            return True
    return False


def caller_list(wt, symbol, windows):
    """The lines of one CALLERS list: hits of the symbol in the worktree outside
    every mapped window of the spec."""
    args = ["grep", "-n", "-I", "--no-color", "-z", "-F"] + (["-w"] if IDENT.match(symbol) else []) + ["-e", symbol]
    raw = git(wt, *args) or b""
    hits = []
    for rec in raw.decode("utf-8", "replace").split("\n"):
        parts = rec.split("\0", 2)
        if len(parts) < 3 or not parts[1].isdigit():
            continue
        path, n = parts[0], int(parts[1])
        if any(path == wp and ws <= n <= we for wp, ws, we in windows):
            continue
        hits.append((path, n, parts[2].rstrip("\r")))
    counts = {}
    for path, _, _ in hits:
        counts[path] = counts.get(path, 0) + 1
    out = ["CALLERS %s: %d hits in %d files" % (symbol, len(hits), len(counts))]
    if len(hits) <= MAX_CALLER_HITS:
        out += ["%s:%d: %s" % (path, n, text[:HIT_TEXT]) for path, n, text in hits]
    else:
        ranked = sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))
        out += ["%6d  %s" % (c, path) for path, c in ranked[:MAX_CALLER_FILES]]
        if len(ranked) > MAX_CALLER_FILES:
            out.append("(%d more files)" % (len(ranked) - MAX_CALLER_FILES))
        out.append("narrow: git grep -n -I %s-F -e %s -- <path>" % ("-w " if IDENT.match(symbol) else "", shlex.quote(symbol)))
    return out


def parse_only(text, p):
    """Labels: <k> for file_changes, r<k> for reads."""
    labels = set()
    for x in text.split(","):
        m = ONLY.match(x)
        if not m:
            p.error("")
        labels.add(m.group(1) + str(int(m.group(2))))
    return labels


def main(argv):
    p = Parser(prog="context-pack.py")
    p.add_argument("spec_id")
    p.add_argument("worktree")
    p.add_argument("--since")
    p.add_argument("--only")
    p.add_argument("--spec-json")
    p.add_argument("--max-lines", type=int, default=DEFAULT_MAX_LINES)
    ns = p.parse_args(argv)
    if ns.max_lines < 1:
        p.error("")
    only = parse_only(ns.only, p) if ns.only is not None else None
    wt = ns.worktree
    if git(wt, "rev-parse", "--git-dir") is None:
        raise Fail(2, "worktree: not a git worktree")
    spec = load_spec(ns.spec_id, ns.spec_json)
    fcs = spec["file_changes"]
    reads = spec.get("reads") or []
    if only is not None:
        if not {x for x in only if x[0] != "r"} <= {str(i) for i in range(1, len(fcs) + 1)}:
            raise Fail(2, "--only: index outside file_changes")
        if not {x for x in only if x[0] == "r"} <= {"r%d" % i for i in range(1, len(reads) + 1)}:
            raise Fail(2, "--only: index outside reads")
    base = str(spec["base"])
    resolved = HEX.match(base) and git(wt, "rev-parse", "--verify", "--quiet", base + "^{commit}")
    tree = Tree(wt, base)
    listed = bool(resolved) and tree.load_changes()
    tree.siblings = lambda: [x for x in (git(wt, "ls-tree", "-r", "--name-only", base) or b"")
                             .decode("utf-8", "replace").split("\n") if x]
    if ns.since is not None and (ns.since.startswith("-") or
                                 git(wt, "rev-parse", "--verify", "--quiet", ns.since + "^{commit}") is None):
        raise Fail(2, "--since: not a commit in this worktree")
    deleted_ok = {d["path"] for d in spec.get("deletes", [])}
    anchors = [resolve(tree, str(i + 1), fc, deleted_ok, bool(resolved), listed) for i, fc in enumerate(fcs)]
    anchors += [resolve(tree, "r%d" % (i + 1), {"location": rd["location"], "change": rd.get("why", "")},
                        deleted_ok, bool(resolved), listed, read=True) for i, rd in enumerate(reads)]

    head = (git(wt, "rev-parse", "--short", "HEAD") or b"none").decode().strip()
    out = ["context-pack %s base %s HEAD %s modified %d" % (
        ns.spec_id, base[:12] if HEX.match(base) else "unresolvable", head, len(tree.changes))]
    out += [table_line(a) for a in sorted(anchors, key=lambda a: a["rank"])]
    out.append("")
    want = [a for a in anchors if a["window"] and (only is None or a["k"] in only)
            and (ns.since is None or overlaps_since(tree, ns.since, a))]
    table_len = len(anchors)
    skipped, skipped_need = [], 0
    left = ns.max_lines - len(out)
    for idx, a in enumerate(want):
        cap, body = window_body(tree, a)
        full = len(body)
        if only is None and full > WINDOW_CAP:
            body = body[:WINDOW_CAP]
        reserve = 1 if idx < len(want) - 1 or skipped else 0   # room for the not-expanded line
        need = 1 + len(body) + (1 if len(body) < full else 0)
        if only is not None and need > left - reserve and left - reserve >= 4:
            body = body[:left - reserve - 2]
            need = 1 + len(body) + 1
        if need > left - reserve:
            skipped.append(a["k"])
            skipped_need += need
            continue
        out.append(cap)
        out += body
        if len(body) < full:
            out.append("%d more lines: --only %s --max-lines %d" % (full - len(body), a["k"], table_len + full + 3))
        left -= need
    if skipped:
        out.append("not expanded: --only " + ",".join(skipped))
        left -= 1
    windows = [(a["window"][1], a["window"][2], a["window"][3]) for a in anchors
               if a["window"] and a["window"][0] == "wt"]
    symbols = list(dict.fromkeys(a["symbol"] for a in want if a["symbol"]))
    lists = [(sym, caller_list(wt, sym, windows)) for sym in symbols[:MAX_CALLER_SYMBOLS]]
    unlisted = symbols[MAX_CALLER_SYMBOLS:]
    cut = []
    fits_all = sum(len(ls) for _, ls in lists) + (1 if unlisted else 0) <= left
    for idx, (sym, ls) in enumerate(lists):
        reserve = 0 if fits_all else 1
        if len(ls) > left - reserve:
            cut.append((sym, len(ls)))
            continue
        out += ls
        left -= len(ls)
    if cut:
        out.append("callers not printed: %s: --max-lines %d" % (
            ",".join(c[0] for c in cut),
            len(out) - (1 if skipped else 0) + skipped_need + sum(c[1] for c in cut) + (1 if unlisted else 0)))
    if unlisted:
        out.append("callers not listed: " + ",".join(unlisted))
    sys.stdout.reconfigure(encoding="utf-8", errors="replace", newline="\n")
    sys.stdout.write("\n".join(out) + "\n")
    return max([a["exit"] for a in anchors] or [0])


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv[1:]))
    except Fail as f:
        sys.stderr.write("context-pack: %s\n" % f.msg)
        sys.exit(f.code)