#!/usr/bin/env python3
"""Conservative worktree lifecycle, serialized with agentctl's worktrees lock.

Git's actual registry is authoritative for removal (including old path names).
Ignored files are deliberately retained: archive evidence and remove disposable
caches explicitly before requesting cleanup; this command never force-removes.
"""
import argparse
from contextlib import contextmanager
import fcntl
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import stat
import subprocess
import sys

# Never leave __pycache__ that dirties an otherwise clean worktree.
sys.dont_write_bytecode = True
os.environ.setdefault('PYTHONDONTWRITEBYTECODE', '1')


class Refused(RuntimeError):
    pass


def maybe_codegraph_init(path, *, timeout=180):
    """Best-effort CodeGraph index for a new worktree. Never blocks create.

    Linked worktrees get their own .codegraph/ when gitignored. Missing CLI or
    unignored path → skip with a notice. Failures are warnings only.
    """
    binary = shutil.which('codegraph')
    if not binary:
        print(json.dumps({'codegraph': 'not_installed'}), file=sys.stderr)
        return {'status': 'not_installed'}
    target = Path(path).resolve()
    ignored = subprocess.run(
        ['git', '-C', str(target), 'check-ignore', '-q', '--', '.codegraph/'],
        capture_output=True, check=False).returncode == 0
    if not ignored:
        print(json.dumps({
            'codegraph': 'skipped',
            'reason': '.codegraph not gitignored',
        }), file=sys.stderr)
        return {'status': 'skipped', 'reason': 'not_ignored'}
    try:
        # Detach so create returns immediately; index can take tens of seconds.
        proc = subprocess.Popen(
            [binary, 'init', '-y', str(target)],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            start_new_session=True,
        )
        print(json.dumps({
            'codegraph': 'init_started',
            'pid': proc.pid,
            'path': str(target),
            'timeout_hint_seconds': timeout,
        }), file=sys.stderr)
        return {'status': 'init_started', 'pid': proc.pid}
    except OSError as exc:
        print(json.dumps({'codegraph': 'init_failed', 'error': str(exc)}), file=sys.stderr)
        return {'status': 'init_failed', 'error': str(exc)}


def command(args, cwd, check=True):
    result = subprocess.run(args, cwd=cwd, text=True, capture_output=True)
    if check and result.returncode:
        raise Refused(result.stderr.strip() or result.stdout.strip() or "command refused")
    return result


def git(repo, *args, check=True):
    return command(["git", "-C", str(repo), *args], repo, check=check)


def worktrees(repo):
    records = []
    for block in git(repo, "worktree", "list", "--porcelain", "-z").stdout.split("\0\0"):
        fields = {}
        for line in block.split("\0"):
            if line:
                key, _, value = line.partition(" ")
                fields[key] = value
        if fields:
            records.append(fields)
    return records


def canonical_main(repo):
    entries = worktrees(repo)
    if not entries or "bare" in entries[0]:
        raise Refused("a non-bare main worktree is required")
    main = Path(entries[0]["worktree"]).resolve()
    common = Path(git(repo, "rev-parse", "--path-format=absolute", "--git-common-dir").stdout.strip()).resolve()
    if common != main / ".git" or not (main / ".git").is_dir():
        raise Refused("main worktree/common Git directory identity mismatch")
    return main


def safe_branch(branch):
    match = re.fullmatch(r"(?:feature|fix)/([A-Za-z0-9][A-Za-z0-9_-]{0,119})", branch)
    if not match:
        raise Refused("branch must be feature/<task-id> or fix/<task-id>; task-id uses letters, digits, - and _")
    return match.group(1)


def safe_directory(main, path):
    parent = main / ".worktrees"
    if parent.is_symlink() or parent.resolve() != parent:
        raise Refused(".worktrees must not be a symlink")
    if path.parent != parent or path.is_symlink() or path.resolve() != path:
        raise Refused("worktree must be a direct, non-symlink child of main/.worktrees")
    if path.name.startswith("."):
        raise Refused("protected worktree path")


def control(repo, *args, check=True):
    # Use the invoked script version but explicitly bind its repository identity.
    return command([sys.executable, str(Path(__file__).with_name("agentctl.py")),
                    "--repo", str(repo), *args], repo, check=check)


@contextmanager
def write_lock(repo, main):
    control(repo, "status")  # Read-only; explicit init is mandatory for writes.
    if os.environ.get("AGENTCTL_RUN_ID") or os.environ.get("AGENTCTL_TASK_ID"):
        control(repo, "guard", "resource", "git:worktrees")
        yield
        return
    lock_dir = main / ".locks" / "os"
    if lock_dir.is_symlink() or lock_dir.resolve() != lock_dir or not lock_dir.is_dir():
        raise Refused("initialized canonical .locks/os is required")
    fd = os.open(lock_dir / "worktrees.lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600)
    try:
        if not stat.S_ISREG(os.fstat(fd).st_mode):
            raise Refused("worktrees lock must be a regular file")
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as exc:
            raise Refused("git:worktrees is busy; retry after the current operation") from exc
        yield
    finally:
        os.close(fd)  # Stable inode: never truncate, replace or unlink lock files.


def create(repo, main, branch):
    slug = safe_branch(branch)
    path = main / ".worktrees" / slug
    safe_directory(main, path)
    entries = worktrees(main)
    if path.exists() or any(Path(e["worktree"]).resolve() == path or
                            e.get("branch") == f"refs/heads/{branch}" for e in entries):
        raise Refused("worktree path or branch is already occupied")
    if git(main, "show-ref", "--verify", "--quiet", f"refs/heads/{branch}", check=False).returncode == 0:
        raise Refused("branch already exists; preserve it and review its worktree mapping")
    if git(main, "check-ignore", "-q", "--", ".worktrees/", check=False).returncode:
        raise Refused("main .gitignore must ignore .worktrees/")
    base = git(main, "rev-parse", "--verify", "refs/heads/dev^{commit}").stdout.strip()
    path.parent.mkdir(exist_ok=True)
    git(main, "worktree", "add", "-b", branch, str(path), base)
    codegraph = maybe_codegraph_init(path)
    print(json.dumps({
        "branch": branch,
        "worktree": str(path),
        "base_dev_sha": base,
        "codegraph": codegraph.get("status"),
    }))


def _linux_same_user_cwds():
    unreadable = []
    for proc in Path("/proc").iterdir():
        if not proc.name.isdigit():
            continue
        try:
            if proc.stat().st_uid != os.getuid():
                continue
            cwd = Path(os.readlink(proc / "cwd"))
            yield proc.name, cwd
        except (FileNotFoundError, ProcessLookupError):
            continue
        except PermissionError:
            # Linux may hide cwd for per-user OS/session helpers despite a matching
            # UID. Skip only when the process is a known session helper, or when
            # cmdline is readable and clearly not rooted under any task worktree.
            # All other unreadable same-user processes remain fail-closed.
            try:
                comm = (proc / "comm").read_text().strip()
            except (OSError, UnicodeError):
                comm = ""
            if comm in {"systemd", "sd-pam", "sshd", "sshd-session"}:
                print(f"[notice] OS session helper {proc.name} ({comm}) hides cwd; "
                      "managed task registry and lifecycle locks remain required", file=sys.stderr)
                continue
            try:
                cmdline = (proc / "cmdline").read_bytes().replace(b"\0", b" ").decode("utf-8", "replace")
            except (OSError, UnicodeError):
                cmdline = ""
            if cmdline and ".worktrees/" not in cmdline and "/.kilo/worktrees/" not in cmdline:
                print(f"[notice] cwd-hidden process {proc.name} ({comm or 'unknown'}) has no worktree path in cmdline; "
                      "managed task registry and lifecycle locks remain required", file=sys.stderr)
                continue
            unreadable.append(proc.name)
    if unreadable:
        raise Refused(f"cannot verify cwd of same-user process {','.join(unreadable)}; cleanup_pending")


def _darwin_same_user_cwds():
    try:
        result = subprocess.run(
            ["lsof", "-a", "-d", "cwd", "-u", str(os.getuid()), "-F", "pn"],
            capture_output=True, text=True)
    except OSError as exc:
        raise Refused(f"cannot enumerate process cwd via lsof ({exc}); cleanup_pending") from exc
    if result.returncode not in (0, 1):
        # lsof exits 1 when the selection is empty; other codes are opaque failure.
        raise Refused(f"cannot enumerate process cwd via lsof (exit {result.returncode}); cleanup_pending")
    if result.stderr.strip() and result.returncode != 0 and not result.stdout.strip():
        raise Refused("cannot verify process cwd via lsof; cleanup_pending")
    pid, saw_cwd = None, False
    for line in result.stdout.splitlines():
        if not line:
            continue
        code, value = line[0], line[1:]
        if code == "p":
            if pid is not None and not saw_cwd:
                raise Refused(f"cannot verify cwd of same-user process {pid}; cleanup_pending")
            pid, saw_cwd = value, False
        elif code == "n" and pid is not None:
            yield pid, Path(value)
            saw_cwd = True
        # macOS lsof may still emit f/type lines; ignore them.
    if pid is not None and not saw_cwd:
        raise Refused(f"cannot verify cwd of same-user process {pid}; cleanup_pending")


def iter_same_user_cwds():
    """Yield (pid_str, cwd Path) for same-uid processes. Fail closed on gaps."""
    if sys.platform.startswith("linux"):
        yield from _linux_same_user_cwds()
    elif sys.platform == "darwin":
        yield from _darwin_same_user_cwds()
    else:
        raise Refused(f"process cwd checks unsupported on {sys.platform}; cleanup_pending")


def assert_no_process_cwd(path):
    path = path.resolve()
    for pid, cwd in iter_same_user_cwds():
        try:
            resolved = cwd.resolve()
        except OSError as exc:
            raise Refused(f"cannot verify cwd of same-user process {pid}; cleanup_pending") from exc
        if resolved == path or path in resolved.parents:
            raise Refused(f"process {pid} still uses this worktree as cwd")


def inspect_process_cwd(path):
    """Structured process-cwd visibility for cleanup reports (never force-clears)."""
    path = Path(path).resolve()
    try:
        for pid, cwd in iter_same_user_cwds():
            try:
                resolved = cwd.resolve()
            except OSError:
                return {'state': 'unreadable', 'unreadable': [str(pid)]}
            if resolved == path or path in resolved.parents:
                return {'state': 'occupied', 'pid': str(pid)}
        return {'state': 'clear'}
    except Refused as exc:
        msg = str(exc)
        if 'unsupported' in msg:
            return {
                'state': 'unsupported_platform',
                'platform': sys.platform,
                'message': msg,
            }
        if 'cannot verify cwd of same-user process' in msg:
            # Extract pid list when present after the word process
            return {'state': 'unreadable', 'unreadable': [msg], 'message': msg}
        return {'state': 'unreadable', 'unreadable': [], 'message': msg}



@contextmanager
def removal_lock(main, path):
    # Every agentctl run/registration holds the shared side of this same lock.
    # Keep it exclusive through removal so no new run can race worktree-check.
    key = hashlib.sha256(str(path.resolve()).encode()).hexdigest()
    lock_path = main / ".locks" / "os" / f"worktree-{key}.lock"
    if lock_path.parent.resolve() != lock_path.parent:
        raise Refused("canonical lifecycle lock directory is required")
    fd = os.open(lock_path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600)
    try:
        if not stat.S_ISREG(os.fstat(fd).st_mode):
            raise Refused("worktree lifecycle lock must be a regular file")
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as exc:
            raise Refused("worktree has an active registered operation; cleanup_pending") from exc
        yield
    finally:
        os.close(fd)


def remove(repo, main, branch, delete_branch):
    safe_branch(branch)
    entries = [entry for entry in worktrees(main) if entry.get("branch") == f"refs/heads/{branch}"]
    if len(entries) != 1:
        raise Refused("expected exactly one registered worktree for the requested branch")
    entry = entries[0]
    path = Path(entry["worktree"])
    safe_directory(main, path)
    if "locked" in entry or "prunable" in entry or not path.is_dir():
        raise Refused("locked, missing or prunable worktree requires manual review")
    with removal_lock(main, path):
        remove_checked(repo, main, branch, delete_branch, path)


def remove_checked(repo, main, branch, delete_branch, path):
    actual = git(path, "symbolic-ref", "--quiet", "HEAD").stdout.strip()
    if actual != f"refs/heads/{branch}":
        raise Refused("worktree branch identity changed")
    head = git(path, "rev-parse", "HEAD").stdout.strip()
    if git(main, "merge-base", "--is-ancestor", head, "refs/heads/dev", check=False).returncode:
        raise Refused("task HEAD is not an ancestor of dev; cleanup_pending")
    for marker in ("MERGE_HEAD", "CHERRY_PICK_HEAD", "REVERT_HEAD", "rebase-merge", "rebase-apply", "BISECT_START"):
        marker_path = Path(git(path, "rev-parse", "--path-format=absolute", "--git-path", marker).stdout.strip())
        if marker_path.exists():
            raise Refused(f"unfinished Git operation: {marker}; cleanup_pending")
    if git(path, "status", "--porcelain=v1", "--untracked-files=all", "--ignored=matching").stdout:
        raise Refused("tracked, untracked or ignored files remain; archive evidence and clean explicitly")
    control(repo, "worktree-check", "--path", str(path))
    assert_no_process_cwd(path)
    if git(path, "rev-parse", "HEAD").stdout.strip() != head:
        raise Refused("worktree HEAD changed during cleanup checks")
    git(main, "worktree", "remove", str(path))
    if delete_branch:
        result = git(main, "branch", "-d", branch, check=False)
        if result.returncode:
            raise Refused(f"worktree removed; branch retained (cleanup_pending): {result.stderr.strip()}")
    print(json.dumps({"removed": str(path), "branch": branch, "branch_deleted": delete_branch}))


def inspect_worktree(repo, main, entry):
    """Read-only cleanup readiness for one porcelain worktree entry."""
    path = Path(entry.get("worktree", ""))
    branch_ref = entry.get("branch", "")
    report = {
        "path": str(path),
        "branch": branch_ref[len("refs/heads/"):] if branch_ref.startswith("refs/heads/") else branch_ref or None,
        "verdict": "blocked",
        "blockers": [],
        "notes": [],
    }
    if "bare" in entry:
        report["blockers"].append("bare_worktree")
        return report
    try:
        safe_directory(main, path.resolve())
    except Refused as exc:
        report["blockers"].append(f"path:{exc}")
        return report
    if "locked" in entry:
        report["blockers"].append("git_locked")
    if "prunable" in entry:
        report["blockers"].append("git_prunable")
    if not path.is_dir():
        report["blockers"].append("missing_directory")
        return report
    if not report["branch"]:
        report["blockers"].append("detached_or_unmapped_branch")
        return report
    try:
        actual = git(path, "symbolic-ref", "--quiet", "HEAD").stdout.strip()
        if actual != f"refs/heads/{report['branch']}":
            report["blockers"].append("branch_identity_mismatch")
        head = git(path, "rev-parse", "HEAD").stdout.strip()
        report["head"] = head
        if git(main, "merge-base", "--is-ancestor", head, "refs/heads/dev", check=False).returncode:
            report["blockers"].append("not_ancestor_of_dev")
        else:
            report["notes"].append("ancestor_of_dev")
        for marker in ("MERGE_HEAD", "CHERRY_PICK_HEAD", "REVERT_HEAD", "rebase-merge", "rebase-apply", "BISECT_START"):
            marker_path = Path(git(path, "rev-parse", "--path-format=absolute", "--git-path", marker).stdout.strip())
            if marker_path.exists():
                report["blockers"].append(f"unfinished_git:{marker}")
        dirty = git(path, "status", "--porcelain=v1", "--untracked-files=all", "--ignored=matching").stdout
        if dirty:
            report["blockers"].append("dirty_or_ignored_files")
            report["dirty_sample"] = dirty.splitlines()[:8]
    except Refused as exc:
        report["blockers"].append(f"git:{exc}")
    check = control(repo, "worktree-check", "--path", str(path), check=False)
    if check.returncode:
        try:
            payload = json.loads(check.stderr.strip().splitlines()[-1])
            report["blockers"].append(f"registry:{payload.get('code', 'DENIED')}")
            if payload.get("message"):
                report["notes"].append(payload["message"])
        except (ValueError, IndexError, json.JSONDecodeError):
            report["blockers"].append("registry:worktree_check_failed")
    else:
        report["notes"].append("registry_clear")
    process = inspect_process_cwd(path)
    report["process"] = process
    if process["state"] == "occupied":
        report["blockers"].append(f"process_cwd:{process['pid']}")
    elif process["state"] == "unreadable":
        report["blockers"].append("process_cwd_unreadable")
    elif process["state"] == "unsupported_platform":
        report["blockers"].append("process_cwd_platform_pending")
        report["notes"].append(process.get("message") or "platform process check pending")
    key = hashlib.sha256(str(path.resolve()).encode()).hexdigest()
    lock_path = main / ".locks" / "os" / f"worktree-{key}.lock"
    if lock_path.exists():
        fd = None
        try:
            fd = os.open(lock_path, os.O_RDWR | os.O_NOFOLLOW, 0o600)
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            fcntl.flock(fd, fcntl.LOCK_UN)
            report["notes"].append("lifecycle_lock_idle")
        except BlockingIOError:
            report["blockers"].append("lifecycle_lock_busy")
        except OSError as exc:
            report["blockers"].append(f"lifecycle_lock:{exc}")
        finally:
            if fd is not None:
                os.close(fd)
    if not report["blockers"]:
        report["verdict"] = "ready"
    return report


def build_cleanup_report(repo, main):
    """Return cleanup readiness payload without deleting anything."""
    try:
        control(repo, "status")
        registry = "initialized"
    except Exception:
        registry = "unavailable"
    items = []
    for entry in worktrees(main):
        path = entry.get("worktree")
        if not path:
            continue
        resolved = Path(path).resolve()
        if resolved == main.resolve():
            continue
        if resolved.parent != (main / ".worktrees").resolve():
            items.append({
                "path": str(resolved),
                "branch": entry.get("branch"),
                "verdict": "blocked",
                "blockers": ["outside_dot_worktrees"],
                "notes": ["not managed by agent-worktree cleanup"],
            })
            continue
        items.append(inspect_worktree(repo, main, entry))
    ready = sum(1 for item in items if item.get("verdict") == "ready")
    return {
        "status": "ok",
        "registry": registry,
        "worktrees": items,
        "summary": {"total": len(items), "ready": ready, "blocked": len(items) - ready},
    }


def cleanup_report(repo, main):
    """Surface cleanup readiness without deleting anything."""
    print(json.dumps(build_cleanup_report(repo, main), ensure_ascii=False, sort_keys=True))

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("list", help="read-only Git worktree listing")
    commands.add_parser("cleanup-report", help="read-only cleanup readiness for .worktrees entries")
    for name in ("create", "remove"):
        child = commands.add_parser(name)
        child.add_argument("branch")
        if name == "remove":
            child.add_argument("--delete-branch", action="store_true")
    args = parser.parse_args()
    repo = Path(__file__).resolve().parent.parent
    try:
        root = canonical_main(repo)
        if args.command == "list":
            print(git(root, "worktree", "list").stdout, end="")
            return 0
        if args.command == "cleanup-report":
            cleanup_report(repo, root)
            return 0
        safe_branch(args.branch)
        with write_lock(repo, root):
            if args.command == "create":
                create(repo, root, args.branch)
            else:
                remove(repo, root, args.branch, args.delete_branch)
        return 0
    except (Refused, OSError) as exc:
        print(f"[refused] {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())