branch_cache.py

"""Per-branch graph cache: data.json (+ optional CRG) keyed by branch@commit."""

from __future__ import annotations

import json
import re
import shutil
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import git_ops
import workspace as workspace_mod

TOOL_DIR = Path(__file__).resolve().parent


def _data_path() -> Path:
    return workspace_mod.get().data_path


def _graphs_dir() -> Path:
    return workspace_mod.get().graphs_dir


def _home() -> Path:
    return workspace_mod.home_dir()


def _utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def _safe_segment(value: str, *, fallback: str = "unknown") -> str:
    s = re.sub(r"[^\w.\-]+", "_", (value or "").strip())[:80]
    return s or fallback


def cache_key(branch: str, commit: str) -> tuple[str, str]:
    short = (commit or "")[:12] or "unknown"
    br = _safe_segment(branch or "detached", fallback="detached")
    return br, short


def cache_dir(branch: str, commit: str) -> Path:
    br, short = cache_key(branch, commit)
    return _graphs_dir() / br / short


def list_caches() -> dict[str, Any]:
    items = []
    graphs = _graphs_dir()
    if not graphs.is_dir():
        return {"ok": True, "caches": items}
    for man in sorted(graphs.rglob("manifest.json"), key=lambda p: p.stat().st_mtime, reverse=True):
        try:
            meta = json.loads(man.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            continue
        items.append(meta)
    return {"ok": True, "caches": items}


def save_current(*, root: Path, crg_status: dict[str, Any] | None = None) -> dict[str, Any]:
    """Snapshot live data.json (+ CRG if present) under graphs/<branch>/<sha>/."""
    data_path = _data_path()
    if not data_path.is_file():
        return {"ok": False, "error": "data.json missing — generate the graph first"}
    ref = git_ops.git_ref(root)
    branch = ref.get("branch") or "detached"
    commit = ref.get("commit") or ""
    if not commit:
        return {"ok": False, "error": "not a git repo / no HEAD"}
    dest = cache_dir(branch, commit)
    dest.mkdir(parents=True, exist_ok=True)
    shutil.copy2(data_path, dest / "data.json")

    crg_src = root / ".code-review-graph"
    crg_copied = False
    crg_error = None
    if crg_src.is_dir():
        try:
            target = dest / "code-review-graph"
            if target.exists():
                shutil.rmtree(target)
            shutil.copytree(crg_src, target)
            crg_copied = True
        except OSError as exc:
            crg_error = str(exc)
    else:
        crg_error = "no .code-review-graph at root"

    try:
        data = json.loads((dest / "data.json").read_text(encoding="utf-8"))
        stats = data.get("stats") or {}
    except (OSError, json.JSONDecodeError):
        stats = {}

    st = crg_status or {}
    manifest = {
        "branch": branch,
        "commit": commit,
        "short": commit[:7],
        "saved_at": _utc_now(),
        "root": str(root),
        "stats": stats,
        "crg": {
            "copied": crg_copied,
            "built_at_commit": st.get("built_at_commit") or st.get("commit"),
            "error": crg_error,
        },
        "path": str(dest.relative_to(_home())).replace("\\", "/"),
    }
    (dest / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    return {"ok": True, "cache": manifest}


def load_into_live(*, root: Path, branch: str | None = None, commit: str | None = None) -> dict[str, Any]:
    """Restore cached data.json (+ CRG) for branch@commit into the live paths."""
    ref = git_ops.git_ref(root)
    br = branch if branch is not None else (ref.get("branch") or "detached")
    sha = commit if commit is not None else (ref.get("commit") or "")
    if not sha:
        return {"ok": False, "error": "no commit", "hit": False}
    dest = cache_dir(br, sha)
    data_file = dest / "data.json"
    if not data_file.is_file():
        return {"ok": False, "error": "cache miss", "hit": False, "branch": br, "commit": sha}
    try:
        shutil.copy2(data_file, _data_path())
    except OSError as exc:
        return {"ok": False, "error": str(exc), "hit": False}

    crg_cached = dest / "code-review-graph"
    crg_restored = False
    if crg_cached.is_dir():
        live = root / ".code-review-graph"
        try:
            if live.exists():
                shutil.rmtree(live)
            shutil.copytree(crg_cached, live)
            crg_restored = True
        except OSError:
            crg_restored = False

    man = {}
    man_path = dest / "manifest.json"
    if man_path.is_file():
        try:
            man = json.loads(man_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            man = {}
    return {
        "ok": True,
        "hit": True,
        "branch": br,
        "commit": sha,
        "crg_restored": crg_restored,
        "cache": man,
    }


def regenerate(*, root: Path, timeout: int = 600) -> dict[str, Any]:
    """Run generate.py --root <root> to rebuild live data.json."""
    script = TOOL_DIR / "generate.py"
    if not script.is_file():
        return {"ok": False, "error": "generate.py missing"}
    try:
        proc = subprocess.run(
            [sys.executable, str(script), "--root", str(root)],
            cwd=str(TOOL_DIR),
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=timeout,
        )
    except (OSError, subprocess.TimeoutExpired) as exc:
        return {"ok": False, "error": str(exc)}
    if proc.returncode != 0:
        return {
            "ok": False,
            "error": (proc.stderr or proc.stdout or "generate failed").strip()[:2000],
        }
    # Cache the freshly generated graph for this HEAD.
    saved = save_current(root=root)
    return {
        "ok": True,
        "stdout": (proc.stdout or "").strip()[-500:],
        "cache": saved.get("cache") if saved.get("ok") else None,
    }


def after_switch(*, root: Path, prefer_cache: bool = True) -> dict[str, Any]:
    """Load cache for new HEAD or regenerate. Call after a successful git switch."""
    ref = git_ops.git_ref(root)
    if prefer_cache:
        loaded = load_into_live(root=root)
        if loaded.get("hit"):
            return {"ok": True, "source": "cache", **loaded, "git": ref}
    gen = regenerate(root=root)
    if not gen.get("ok"):
        return {"ok": False, "source": "generate", "error": gen.get("error"), "git": ref}
    return {"ok": True, "source": "generate", "git": ref, "cache": gen.get("cache")}