#!/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>[,<k>...]]
                       [--spec-json <file>] [--max-lines <n>]

Prints a header, a status table (one line per anchor) and bounded windows of the
code each anchor names, numbered in the worktree's lines. Exit: 0 every anchor
carried; 3 an anchor cannot be carried (OUT_OF_RANGE, UNRESOLVED, 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, not the table: the table always lists every
anchor, so a budget smaller than the table prints the table and no window.
"""
import argparse
import json
import os
import re
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+))? @@")
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")
    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)
        loc = fc.get("location")
        if not isinstance(loc, dict):
            bad(at + ".location")
        if not isinstance(loc.get("path"), str):
            bad(at + ".location.path")
        if not is_int(loc.get("line")):
            bad(at + ".location.line")
        if loc.get("end_line") is not None and not is_int(loc["end_line"]):
            bad(at + ".location.end_line")
        if "change" in fc and not isinstance(fc["change"], str):
            bad(at + ".change")
    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 resolve(tree, k, fc, deleted_ok, base_ok, listed):
    loc = fc["location"]
    path, s = loc["path"], int(loc["line"])
    e = int(loc.get("end_line") or s)
    a = {"k": k, "path": path, "s": s, "e": e, "change": str(fc.get("change", "")),
         "rank": 1, "exit": 0, "status": "", "now": None, "window": None}
    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:
        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 = "[%d] %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)
    cap = "== [%d] %s:%d-%d%s  %s" % (a["k"], path, s, e, " AT BASE" if kind == "base" else "", a["status"])
    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 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 = None
    if ns.only is not None:
        try:
            only = {int(x) for x in ns.only.split(",")}
        except ValueError:
            p.error("")
    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"]
    if only is not None and not only <= set(range(1, len(fcs) + 1)):
        raise Fail(2, "--only: index outside file_changes")
    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, i + 1, fc, deleted_ok, bool(resolved), listed) for i, fc in enumerate(fcs)]

    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(fcs)
    skipped = []
    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"])
            continue
        out.append(cap)
        out += body
        if len(body) < full:
            out.append("%d more lines: --only %d --max-lines %d" % (full - len(body), a["k"], table_len + full + 3))
        left -= need
    if skipped:
        out.append("not expanded: --only " + ",".join(map(str, skipped)))
    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)