"""Pre-edit snapshots, wired as an assistant PreToolUse hook.

The launcher registers ``python -c "from zytectx.undo_shield import main; main()"``
against Write|Edit|Bash|NotebookEdit. The hook contract: tool payload as JSON on
stdin, exit 0 to allow the tool call. This hook only observes - it snapshots the
target file before it changes and must NEVER block, so every failure path still
exits 0. A crash here would take down every tool call the assistant makes.

    python -m zytectx.undo_shield --list           recent snapshots
    python -m zytectx.undo_shield --restore PATH   restore the newest snapshot of PATH
"""

from __future__ import annotations

import hashlib
import json
import os
import shutil
import sys
import time
from pathlib import Path

SNAPSHOT_TOOLS = {"Write", "Edit", "NotebookEdit"}
MAX_SNAPSHOT_BYTES = 512 * 1024
MAX_ENTRIES = 200
MAX_TOTAL_BYTES = 50 * 1024 * 1024


def _undo_dir() -> Path:
    value = os.environ.get("ZYTE_DATA_DIR", "").strip()
    base = Path(value) if value else Path.cwd() / ".zyte"
    return base / "undo"


def _index_path() -> Path:
    return _undo_dir() / "index.jsonl"


def _read_index() -> list[dict]:
    entries: list[dict] = []
    try:
        with open(_index_path(), encoding="utf-8") as handle:
            for line in handle:
                try:
                    entry = json.loads(line)
                except ValueError:
                    continue
                if isinstance(entry, dict):
                    entries.append(entry)
    except OSError:
        pass
    return entries


def _write_index(entries: list[dict]) -> None:
    _undo_dir().mkdir(parents=True, exist_ok=True)
    with open(_index_path(), "w", encoding="utf-8") as handle:
        for entry in entries:
            handle.write(json.dumps(entry) + "\n")


def _prune(entries: list[dict]) -> list[dict]:
    kept = entries[-MAX_ENTRIES:]
    total = 0
    survivors: list[dict] = []
    for entry in reversed(kept):
        total += int(entry.get("bytes") or 0)
        if total > MAX_TOTAL_BYTES:
            break
        survivors.append(entry)
    survivors.reverse()
    for entry in entries:
        if entry in survivors:
            continue
        snapshot = entry.get("snapshot")
        if isinstance(snapshot, str):
            shutil.rmtree(Path(snapshot).parent, ignore_errors=True)
    return survivors


def snapshot(payload: dict) -> dict | None:
    """Copy the tool's target file aside. Returns the index entry, or None."""
    tool = payload.get("tool_name")
    if tool not in SNAPSHOT_TOOLS:
        return None
    tool_input = payload.get("tool_input")
    target = tool_input.get("file_path") if isinstance(tool_input, dict) else None
    if not isinstance(target, str) or not target:
        return None
    source = Path(target)
    try:
        if not source.is_file() or source.stat().st_size > MAX_SNAPSHOT_BYTES:
            return None
        stamp = int(time.time() * 1000)
        digest = hashlib.sha256(str(source).encode("utf-8", errors="replace")).hexdigest()[:8]
        folder = _undo_dir() / f"{stamp}-{digest}"
        folder.mkdir(parents=True, exist_ok=True)
        copy = folder / source.name
        shutil.copyfile(source, copy)
        entry = {
            "ts": stamp / 1000,
            "tool": tool,
            "path": str(source),
            "snapshot": str(copy),
            "bytes": copy.stat().st_size,
        }
        _write_index(_prune([*_read_index(), entry]))
        return entry
    except OSError:
        return None


def _list_snapshots() -> int:
    entries = _read_index()
    if not entries:
        print("No snapshots.")
        return 0
    for entry in entries[-20:]:
        when = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(entry.get("ts", 0)))
        print(f"{when}  {entry.get('tool', '?'):<12} {entry.get('path', '?')}")
    return 0


def _restore(path: str | None) -> int:
    entries = _read_index()
    if path:
        wanted = str(Path(path))
        entries = [entry for entry in entries if entry.get("path") == wanted]
    if not entries:
        print("Nothing to restore.", file=sys.stderr)
        return 1
    entry = entries[-1]
    source = Path(str(entry.get("snapshot")))
    target = Path(str(entry.get("path")))
    if not source.is_file():
        print(f"Snapshot missing: {source}", file=sys.stderr)
        return 1
    target.parent.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(source, target)
    print(f"Restored {target} from {source}")
    return 0


def main(argv: list[str] | None = None) -> int:
    arguments = sys.argv[1:] if argv is None else argv
    if "--list" in arguments:
        return _list_snapshots()
    if "--restore" in arguments:
        index = arguments.index("--restore")
        return _restore(arguments[index + 1] if index + 1 < len(arguments) else None)
    # Hook mode: observe-only, so every path out of here is exit 0.
    try:
        payload = json.loads(sys.stdin.read() or "{}")
        if isinstance(payload, dict):
            snapshot(payload)
    except Exception:
        pass
    return 0


if __name__ == "__main__":
    raise SystemExit(main())