#!/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 stat
import subprocess
import sys


class Refused(RuntimeError):
    pass


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)
    print(json.dumps({"branch": branch, "worktree": str(path), "base_dev_sha": base}))


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 the per-user OS service manager despite a
            # matching UID. This narrowly named OS process is not a task agent;
            # all other unreadable same-user processes remain fail-closed.
            try:
                if (proc / "comm").read_text().strip() == "systemd":
                    print(f"[notice] OS user service manager {proc.name} hides cwd; "
                          "managed task registry and lifecycle locks remain required", file=sys.stderr)
                    continue
            except (OSError, UnicodeError):
                pass
            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")


@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 main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("list", help="read-only Git worktree listing")
    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
        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())