Git Worktree Manager

#!/usr/bin/env python3
"""
port_allocator.py - Manage port allocation across git worktrees.

Assigns deterministic port blocks to worktrees, detects conflicts with
running processes, and maintains a central port registry to prevent collisions.

Usage:
    python port_allocator.py status
    python port_allocator.py assign wt-auth --index 1
    python port_allocator.py release wt-auth
    python port_allocator.py check --port 3010
    python port_allocator.py status --json
"""

import argparse
import json
import os
import socket
import subprocess
import sys
import time
from pathlib import Path


DEFAULT_BASE_PORTS = {
    "app": 3000,
    "database": 5432,
    "redis": 6379,
    "api": 8000,
}
DEFAULT_STRIDE = 10
REGISTRY_FILE = ".worktree-ports-registry.json"


def run_git(args, cwd=None):
    """Run a git command and return stdout, or None on failure."""
    try:
        result = subprocess.run(
            ["git"] + args,
            capture_output=True, text=True, cwd=cwd, timeout=30
        )
        return result.stdout.strip() if result.returncode == 0 else None
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return None


def get_repo_root():
    """Get the main repository root."""
    root = run_git(["rev-parse", "--show-toplevel"])
    if not root:
        print("Error: not inside a git repository.", file=sys.stderr)
        sys.exit(1)
    return root


def registry_path(repo_root):
    """Path to the central port registry file."""
    return os.path.join(repo_root, REGISTRY_FILE)


def load_registry(repo_root):
    """Load the port registry from disk."""
    path = registry_path(repo_root)
    if not os.path.isfile(path):
        return {"version": 1, "stride": DEFAULT_STRIDE, "base_ports": DEFAULT_BASE_PORTS, "allocations": {}}
    try:
        with open(path, "r") as f:
            data = json.load(f)
        # Ensure required keys
        data.setdefault("version", 1)
        data.setdefault("stride", DEFAULT_STRIDE)
        data.setdefault("base_ports", DEFAULT_BASE_PORTS)
        data.setdefault("allocations", {})
        return data
    except (json.JSONDecodeError, OSError) as e:
        print(f"Warning: could not read registry ({e}), starting fresh.", file=sys.stderr)
        return {"version": 1, "stride": DEFAULT_STRIDE, "base_ports": DEFAULT_BASE_PORTS, "allocations": {}}


def save_registry(repo_root, registry):
    """Persist the port registry to disk."""
    path = registry_path(repo_root)
    try:
        with open(path, "w") as f:
            json.dump(registry, f, indent=2)
            f.write("\n")
    except OSError as e:
        print(f"Error: could not write registry: {e}", file=sys.stderr)
        sys.exit(1)


def is_port_in_use(port):
    """Check if a TCP port is currently in use by attempting to bind."""
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(1)
            result = s.connect_ex(("127.0.0.1", port))
            return result == 0
    except OSError:
        return False


def compute_ports(index, base_ports, stride):
    """Compute the port block for a given worktree index."""
    return {
        service: base + (index * stride)
        for service, base in base_ports.items()
    }


def find_next_index(registry):
    """Find the lowest available index not in use."""
    used_indices = set()
    for alloc in registry["allocations"].values():
        used_indices.add(alloc.get("index", 0))
    idx = 0
    while idx in used_indices:
        idx += 1
    return idx


def get_active_worktrees(repo_root):
    """Get set of worktree directory names from git."""
    raw = run_git(["worktree", "list", "--porcelain"], cwd=repo_root)
    if not raw:
        return set()
    names = set()
    for line in raw.splitlines():
        if line.startswith("worktree "):
            path = line[len("worktree "):]
            names.add(os.path.basename(path))
    return names


def cmd_status(args):
    """Show current port allocations and their status."""
    repo_root = get_repo_root()
    registry = load_registry(repo_root)
    allocations = registry["allocations"]
    active_wts = get_active_worktrees(repo_root)

    entries = []
    for name, alloc in sorted(allocations.items()):
        ports = alloc.get("ports", {})
        port_status = {}
        for service, port in ports.items():
            port_status[service] = {
                "port": port,
                "in_use": is_port_in_use(port),
            }
        entries.append({
            "worktree": name,
            "index": alloc.get("index", "?"),
            "branch": alloc.get("branch", "?"),
            "ports": port_status,
            "active": name in active_wts or os.path.basename(repo_root) == name,
            "created": alloc.get("created", "?"),
        })

    if args.json:
        output = {
            "repo_root": repo_root,
            "stride": registry["stride"],
            "base_ports": registry["base_ports"],
            "allocation_count": len(entries),
            "allocations": entries,
        }
        print(json.dumps(output, indent=2))
        return

    print(f"Port Allocation Status")
    print(f"Repository: {repo_root}")
    print(f"Stride: {registry['stride']}  |  Base ports: {', '.join(f'{s}={p}' for s, p in registry['base_ports'].items())}")
    print("=" * 90)

    if not entries:
        print("\nNo port allocations registered.")
        print("Use 'assign <worktree-name>' to allocate ports.")
        return

    print(f"\n{'Worktree':<25} {'Idx':<5} {'Branch':<25} {'Ports':<30} {'Status'}")
    print("-" * 90)

    for e in entries:
        ports_str = ", ".join(
            f"{s}:{p['port']}" for s, p in e["ports"].items()
        )
        busy = [s for s, p in e["ports"].items() if p["in_use"]]
        if not e["active"]:
            status = "ORPHANED"
        elif busy:
            status = f"BUSY({','.join(busy)})"
        else:
            status = "free"

        wt_display = e["worktree"] if len(e["worktree"]) <= 24 else e["worktree"][:21] + "..."
        br_display = e["branch"] if len(e["branch"]) <= 24 else e["branch"][:21] + "..."
        print(f"{wt_display:<25} {e['index']:<5} {br_display:<25} {ports_str:<30} {status}")

    total_ports = sum(len(e["ports"]) for e in entries)
    busy_total = sum(1 for e in entries for p in e["ports"].values() if p["in_use"])
    orphaned = sum(1 for e in entries if not e["active"])
    print(f"\nTotal: {len(entries)} allocation(s), {total_ports} port(s), {busy_total} in use, {orphaned} orphaned")


def cmd_assign(args):
    """Assign a port block to a worktree."""
    repo_root = get_repo_root()
    registry = load_registry(repo_root)
    name = args.name

    if name in registry["allocations"]:
        existing = registry["allocations"][name]
        if args.json:
            print(json.dumps({"error": f"'{name}' already allocated", "existing": existing}))
        else:
            print(f"Error: '{name}' already has an allocation (index {existing.get('index')}).", file=sys.stderr)
            print("Use 'release' first, then re-assign.", file=sys.stderr)
        sys.exit(1)

    index = args.index if args.index is not None else find_next_index(registry)
    base_ports = registry["base_ports"]
    stride = registry["stride"]
    ports = compute_ports(index, base_ports, stride)

    # Check for index collision
    for existing_name, alloc in registry["allocations"].items():
        if alloc.get("index") == index:
            msg = f"Index {index} is already assigned to '{existing_name}'."
            if args.json:
                print(json.dumps({"error": msg}))
            else:
                print(f"Error: {msg}", file=sys.stderr)
            sys.exit(1)

    # Check for port conflicts with running processes
    conflicts = []
    for service, port in ports.items():
        if is_port_in_use(port):
            conflicts.append({"service": service, "port": port})

    allocation = {
        "index": index,
        "branch": args.branch or "",
        "ports": ports,
        "created": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }

    registry["allocations"][name] = allocation
    save_registry(repo_root, registry)

    # Also write per-worktree port file if worktree exists
    parent_dir = Path(repo_root).parent
    wt_path = parent_dir / name
    if wt_path.is_dir():
        wt_port_file = wt_path / ".worktree-ports.json"
        try:
            with open(wt_port_file, "w") as f:
                json.dump({
                    "worktree": name,
                    "branch": args.branch or "",
                    "index": index,
                    "ports": ports,
                    "created": allocation["created"],
                }, f, indent=2)
                f.write("\n")
        except OSError:
            pass  # Non-fatal: registry is the source of truth

    if args.json:
        output = {
            "action": "assigned",
            "worktree": name,
            "index": index,
            "ports": ports,
            "conflicts": conflicts,
        }
        print(json.dumps(output, indent=2))
        return

    print(f"Port block assigned to '{name}' (index {index}):\n")
    for service, port in sorted(ports.items()):
        busy = " (IN USE)" if any(c["port"] == port for c in conflicts) else ""
        print(f"  {service:<12} {port}{busy}")

    if conflicts:
        print(f"\nWarning: {len(conflicts)} port(s) currently in use by other processes.")

    print(f"\nRegistry saved to {registry_path(repo_root)}")


def cmd_release(args):
    """Release a port allocation for a worktree."""
    repo_root = get_repo_root()
    registry = load_registry(repo_root)
    name = args.name

    if name not in registry["allocations"]:
        msg = f"No allocation found for '{name}'."
        if args.json:
            print(json.dumps({"error": msg}))
        else:
            print(f"Error: {msg}", file=sys.stderr)
        sys.exit(1)

    released = registry["allocations"].pop(name)
    save_registry(repo_root, registry)

    if args.json:
        print(json.dumps({"action": "released", "worktree": name, "released": released}))
    else:
        print(f"Released port allocation for '{name}' (index {released.get('index')}).")
        print(f"Freed ports: {', '.join(f'{s}={p}' for s, p in released.get('ports', {}).items())}")


def cmd_check(args):
    """Check if specific ports are available."""
    ports_to_check = [int(p.strip()) for p in args.port.split(",")]
    repo_root = get_repo_root()
    registry = load_registry(repo_root)

    # Build reverse map: port -> (worktree, service)
    port_owners = {}
    for name, alloc in registry["allocations"].items():
        for service, port in alloc.get("ports", {}).items():
            port_owners[port] = {"worktree": name, "service": service}

    results = []
    for port in ports_to_check:
        in_use = is_port_in_use(port)
        owner = port_owners.get(port)
        results.append({
            "port": port,
            "in_use": in_use,
            "allocated_to": owner,
            "available": not in_use and owner is None,
        })

    if args.json:
        print(json.dumps({"ports": results}, indent=2))
        return

    print(f"{'Port':<8} {'Process':<12} {'Allocated To':<30} {'Available'}")
    print("-" * 60)
    for r in results:
        proc = "BUSY" if r["in_use"] else "free"
        owner = ""
        if r["allocated_to"]:
            owner = f"{r['allocated_to']['worktree']} ({r['allocated_to']['service']})"
        avail = "yes" if r["available"] else "NO"
        print(f"{r['port']:<8} {proc:<12} {owner:<30} {avail}")


def cmd_sync(args):
    """Sync registry with actual git worktrees, removing orphaned entries."""
    repo_root = get_repo_root()
    registry = load_registry(repo_root)
    active_wts = get_active_worktrees(repo_root)

    orphaned = []
    for name in list(registry["allocations"].keys()):
        if name not in active_wts and os.path.basename(repo_root) != name:
            orphaned.append(name)

    if not orphaned:
        if args.json:
            print(json.dumps({"action": "sync", "removed": [], "message": "Registry is in sync."}))
        else:
            print("Registry is in sync with active worktrees. Nothing to clean.")
        return

    if args.dry_run:
        if args.json:
            print(json.dumps({"action": "sync_preview", "orphaned": orphaned}))
        else:
            print(f"Orphaned allocations ({len(orphaned)}):")
            for name in orphaned:
                idx = registry["allocations"][name].get("index", "?")
                print(f"  {name} (index {idx})")
            print("\nRe-run without --dry-run to remove these entries.")
        return

    removed = []
    for name in orphaned:
        registry["allocations"].pop(name)
        removed.append(name)

    save_registry(repo_root, registry)

    if args.json:
        print(json.dumps({"action": "sync", "removed": removed}))
    else:
        print(f"Removed {len(removed)} orphaned allocation(s):")
        for name in removed:
            print(f"  {name}")
        print(f"\nRegistry saved to {registry_path(repo_root)}")


def main():
    parser = argparse.ArgumentParser(
        description="Manage port allocation across git worktrees to prevent conflicts."
    )
    parser.add_argument("--json", action="store_true", help="Output in JSON format")
    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # status
    sp_status = subparsers.add_parser("status", help="Show all port allocations and their status")
    sp_status.add_argument("--json", action="store_true", help="Output in JSON format")

    # assign
    sp_assign = subparsers.add_parser("assign", help="Assign a port block to a worktree")
    sp_assign.add_argument("name", help="Worktree name (e.g., wt-auth)")
    sp_assign.add_argument("--index", type=int, default=None, help="Worktree index (auto-assigned if omitted)")
    sp_assign.add_argument("--branch", type=str, default="", help="Branch name for metadata")
    sp_assign.add_argument("--json", action="store_true", help="Output in JSON format")

    # release
    sp_release = subparsers.add_parser("release", help="Release a port allocation")
    sp_release.add_argument("name", help="Worktree name to release")
    sp_release.add_argument("--json", action="store_true", help="Output in JSON format")

    # check
    sp_check = subparsers.add_parser("check", help="Check if ports are available")
    sp_check.add_argument("--port", required=True, help="Comma-separated port(s) to check")
    sp_check.add_argument("--json", action="store_true", help="Output in JSON format")

    # sync
    sp_sync = subparsers.add_parser("sync", help="Sync registry with active worktrees")
    sp_sync.add_argument("--dry-run", action="store_true", help="Preview without removing")
    sp_sync.add_argument("--json", action="store_true", help="Output in JSON format")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    commands = {
        "status": cmd_status,
        "assign": cmd_assign,
        "release": cmd_release,
        "check": cmd_check,
        "sync": cmd_sync,
    }
    commands[args.command](args)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
worktree_manager.py - List, create, and clean up git worktrees with status information.

Reports branch, dirty state, age, and merge status for each worktree.
Supports creating new worktrees with deterministic naming and cleaning up
stale or merged-branch worktrees safely.

Usage:
    python worktree_manager.py list
    python worktree_manager.py list --json
    python worktree_manager.py create feature/auth --base main
    python worktree_manager.py remove ../wt-auth
    python worktree_manager.py cleanup --stale-days 14 --dry-run
"""

import argparse
import json
import os
import re
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path


def run_git(args, cwd=None):
    """Run a git command and return stdout, or None on failure."""
    try:
        result = subprocess.run(
            ["git"] + args,
            capture_output=True, text=True, cwd=cwd, timeout=30
        )
        if result.returncode == 0:
            return result.stdout.strip()
        return None
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return None


def get_repo_root():
    """Get the top-level directory of the main git repository."""
    root = run_git(["rev-parse", "--show-toplevel"])
    if not root:
        print("Error: not inside a git repository.", file=sys.stderr)
        sys.exit(1)
    return root


def parse_worktree_list_porcelain(repo_root):
    """Parse 'git worktree list --porcelain' into structured data."""
    raw = run_git(["worktree", "list", "--porcelain"], cwd=repo_root)
    if raw is None:
        return []

    worktrees = []
    current = {}
    for line in raw.splitlines():
        if line.startswith("worktree "):
            if current:
                worktrees.append(current)
            current = {"path": line[len("worktree "):]}
        elif line.startswith("HEAD "):
            current["head"] = line[len("HEAD "):]
        elif line.startswith("branch "):
            ref = line[len("branch "):]
            current["branch"] = ref.replace("refs/heads/", "")
        elif line == "detached":
            current["branch"] = "(detached HEAD)"
            current["detached"] = True
        elif line == "bare":
            current["bare"] = True

    if current:
        worktrees.append(current)
    return worktrees


def get_worktree_status(wt_path):
    """Check dirty state of a worktree."""
    porcelain = run_git(["status", "--porcelain"], cwd=wt_path)
    if porcelain is None:
        return "unknown"
    return "dirty" if porcelain else "clean"


def get_worktree_age_days(wt_path):
    """Get the age of a worktree directory in days."""
    try:
        mtime = os.path.getmtime(wt_path)
        age_seconds = time.time() - mtime
        return max(0, int(age_seconds / 86400))
    except OSError:
        return -1


def is_branch_merged(branch, target="main", cwd=None):
    """Check whether a branch has been merged into target."""
    merged = run_git(["branch", "--merged", target], cwd=cwd)
    if merged is None:
        return False
    merged_branches = [b.strip().lstrip("* ") for b in merged.splitlines()]
    return branch in merged_branches


def get_ahead_behind(branch, cwd=None):
    """Get ahead/behind counts relative to the upstream or origin/main."""
    upstream = run_git(["rev-parse", "--abbrev-ref", f"{branch}@{{upstream}}"], cwd=cwd)
    if not upstream:
        upstream = "origin/main"
    counts = run_git(["rev-list", "--left-right", "--count", f"{branch}...{upstream}"], cwd=cwd)
    if counts:
        parts = counts.split()
        if len(parts) == 2:
            return int(parts[0]), int(parts[1])
    return 0, 0


def enrich_worktree(wt, repo_root):
    """Add status, age, merge info, and ahead/behind to a worktree dict."""
    path = wt.get("path", "")
    is_main = (path == repo_root)
    wt["is_main"] = is_main
    wt["exists"] = os.path.isdir(path)

    if wt["exists"]:
        wt["status"] = get_worktree_status(path)
        wt["age_days"] = get_worktree_age_days(path)
    else:
        wt["status"] = "missing"
        wt["age_days"] = -1

    branch = wt.get("branch", "")
    if branch and branch != "(detached HEAD)" and not is_main:
        wt["merged"] = is_branch_merged(branch, cwd=repo_root)
        ahead, behind = get_ahead_behind(branch, cwd=repo_root)
        wt["ahead"] = ahead
        wt["behind"] = behind
    else:
        wt["merged"] = False
        wt["ahead"] = 0
        wt["behind"] = 0

    return wt


def cmd_list(args):
    """List all worktrees with enriched status."""
    repo_root = get_repo_root()
    worktrees = parse_worktree_list_porcelain(repo_root)
    enriched = [enrich_worktree(wt, repo_root) for wt in worktrees]

    if args.json:
        print(json.dumps(enriched, indent=2))
        return

    if not enriched:
        print("No worktrees found.")
        return

    print(f"{'Path':<45} {'Branch':<30} {'Status':<10} {'Age':<8} {'Flags'}")
    print("-" * 110)
    for wt in enriched:
        path = wt.get("path", "?")
        branch = wt.get("branch", "?")
        status = wt.get("status", "?")
        age = f"{wt['age_days']}d" if wt["age_days"] >= 0 else "?"

        flags = []
        if wt.get("is_main"):
            flags.append("MAIN")
        if wt.get("merged"):
            flags.append("MERGED")
        if wt.get("detached"):
            flags.append("DETACHED")
        if not wt.get("exists"):
            flags.append("MISSING")
        ahead, behind = wt.get("ahead", 0), wt.get("behind", 0)
        if ahead:
            flags.append(f"+{ahead}")
        if behind:
            flags.append(f"-{behind}")

        flag_str = ", ".join(flags) if flags else ""
        # Truncate long paths/branches for display
        display_path = path if len(path) <= 44 else "..." + path[-41:]
        display_branch = branch if len(branch) <= 29 else branch[:26] + "..."
        print(f"{display_path:<45} {display_branch:<30} {status:<10} {age:<8} {flag_str}")

    total = len(enriched)
    dirty = sum(1 for w in enriched if w.get("status") == "dirty")
    merged = sum(1 for w in enriched if w.get("merged"))
    print(f"\nTotal: {total} | Dirty: {dirty} | Merged: {merged}")


def cmd_create(args):
    """Create a new worktree with deterministic naming."""
    repo_root = get_repo_root()
    branch = args.branch
    base = args.base

    # Derive worktree name from branch
    short_name = re.sub(r".*/", "", branch).lower()
    short_name = re.sub(r"[^a-z0-9-]", "-", short_name)
    wt_name = f"wt-{short_name}"
    parent_dir = Path(repo_root).parent
    wt_path = str(parent_dir / wt_name)

    if os.path.exists(wt_path):
        msg = f"Error: path already exists: {wt_path}"
        if args.json:
            print(json.dumps({"error": msg}))
        else:
            print(msg, file=sys.stderr)
        sys.exit(1)

    # Check if branch already exists
    branch_exists = run_git(["rev-parse", "--verify", branch], cwd=repo_root) is not None

    if branch_exists:
        result = run_git(["worktree", "add", wt_path, branch], cwd=repo_root)
    else:
        result = run_git(["worktree", "add", wt_path, "-b", branch, base], cwd=repo_root)

    if result is None:
        msg = f"Error: failed to create worktree at {wt_path}"
        if args.json:
            print(json.dumps({"error": msg}))
        else:
            print(msg, file=sys.stderr)
        sys.exit(1)

    info = {
        "action": "created",
        "path": wt_path,
        "branch": branch,
        "base": base,
        "name": wt_name,
    }

    if args.json:
        print(json.dumps(info, indent=2))
    else:
        print(f"Worktree created: {wt_path}")
        print(f"  Branch: {branch}")
        print(f"  Base:   {base}")
        print(f"  Name:   {wt_name}")
        print(f"\nNext: cd {wt_path}")


def cmd_remove(args):
    """Remove a worktree safely."""
    repo_root = get_repo_root()
    wt_path = os.path.abspath(args.path)

    if wt_path == repo_root:
        msg = "Error: cannot remove the main worktree."
        if args.json:
            print(json.dumps({"error": msg}))
        else:
            print(msg, file=sys.stderr)
        sys.exit(1)

    # Check dirty state unless --force
    if not args.force and os.path.isdir(wt_path):
        status = get_worktree_status(wt_path)
        if status == "dirty":
            msg = f"Worktree at {wt_path} has uncommitted changes. Use --force to remove."
            if args.json:
                print(json.dumps({"error": msg, "status": "dirty"}))
            else:
                print(msg, file=sys.stderr)
            sys.exit(1)

    git_args = ["worktree", "remove", wt_path]
    if args.force:
        git_args.insert(2, "--force")

    result = run_git(git_args, cwd=repo_root)
    if result is None:
        msg = f"Error: failed to remove worktree at {wt_path}"
        if args.json:
            print(json.dumps({"error": msg}))
        else:
            print(msg, file=sys.stderr)
        sys.exit(1)

    info = {"action": "removed", "path": wt_path, "forced": args.force}
    if args.json:
        print(json.dumps(info, indent=2))
    else:
        print(f"Worktree removed: {wt_path}")


def cmd_cleanup(args):
    """Clean up stale and merged-branch worktrees."""
    repo_root = get_repo_root()
    worktrees = parse_worktree_list_porcelain(repo_root)
    enriched = [enrich_worktree(wt, repo_root) for wt in worktrees]

    candidates = []
    for wt in enriched:
        if wt.get("is_main"):
            continue
        reasons = []
        if wt.get("merged"):
            reasons.append("merged")
        if wt.get("age_days", 0) >= args.stale_days:
            reasons.append(f"stale ({wt['age_days']}d)")
        if not wt.get("exists"):
            reasons.append("missing")
        if reasons:
            wt["cleanup_reasons"] = reasons
            wt["safe_to_remove"] = wt.get("status") != "dirty"
            candidates.append(wt)

    if args.json:
        output = {
            "dry_run": args.dry_run,
            "stale_threshold_days": args.stale_days,
            "candidates": candidates,
        }
        if not args.dry_run:
            removed = []
            skipped = []
            for c in candidates:
                if c["safe_to_remove"]:
                    run_git(["worktree", "remove", c["path"]], cwd=repo_root)
                    removed.append(c["path"])
                else:
                    skipped.append(c["path"])
            run_git(["worktree", "prune"], cwd=repo_root)
            output["removed"] = removed
            output["skipped_dirty"] = skipped
        print(json.dumps(output, indent=2))
        return

    if not candidates:
        print(f"No cleanup candidates (threshold: {args.stale_days} days).")
        run_git(["worktree", "prune"], cwd=repo_root)
        return

    print(f"Cleanup candidates (threshold: {args.stale_days} days):\n")
    for c in candidates:
        reasons = ", ".join(c["cleanup_reasons"])
        safe = "safe" if c["safe_to_remove"] else "DIRTY - skip"
        print(f"  {c['path']}")
        print(f"    Branch:  {c.get('branch', '?')}")
        print(f"    Reasons: {reasons}")
        print(f"    Status:  {safe}")
        print()

    if args.dry_run:
        print(f"Dry run: {len(candidates)} candidate(s). Re-run without --dry-run to remove.")
    else:
        removed = 0
        for c in candidates:
            if c["safe_to_remove"]:
                run_git(["worktree", "remove", c["path"]], cwd=repo_root)
                print(f"  Removed: {c['path']}")
                removed += 1
            else:
                print(f"  Skipped (dirty): {c['path']}")
        run_git(["worktree", "prune"], cwd=repo_root)
        print(f"\nRemoved {removed} worktree(s). Pruned stale metadata.")


def main():
    parser = argparse.ArgumentParser(
        description="Manage git worktrees: list, create, remove, and cleanup."
    )
    parser.add_argument(
        "--json", action="store_true", help="Output in JSON format"
    )
    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # list
    sp_list = subparsers.add_parser("list", help="List all worktrees with status")
    sp_list.add_argument("--json", action="store_true", help="Output in JSON format")

    # create
    sp_create = subparsers.add_parser("create", help="Create a new worktree")
    sp_create.add_argument("branch", help="Branch name for the worktree")
    sp_create.add_argument("--base", default="main", help="Base branch (default: main)")
    sp_create.add_argument("--json", action="store_true", help="Output in JSON format")

    # remove
    sp_remove = subparsers.add_parser("remove", help="Remove a worktree")
    sp_remove.add_argument("path", help="Path to the worktree to remove")
    sp_remove.add_argument("--force", action="store_true", help="Force removal even if dirty")
    sp_remove.add_argument("--json", action="store_true", help="Output in JSON format")

    # cleanup
    sp_cleanup = subparsers.add_parser("cleanup", help="Clean up stale/merged worktrees")
    sp_cleanup.add_argument(
        "--stale-days", type=int, default=14,
        help="Consider worktrees older than N days stale (default: 14)"
    )
    sp_cleanup.add_argument(
        "--dry-run", action="store_true",
        help="Show candidates without removing"
    )
    sp_cleanup.add_argument("--json", action="store_true", help="Output in JSON format")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    commands = {
        "list": cmd_list,
        "create": cmd_create,
        "remove": cmd_remove,
        "cleanup": cmd_cleanup,
    }
    commands[args.command](args)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
worktree_validator.py - Validate worktree health across a git repository.

Checks for stale worktrees, missing directories, orphaned branches,
environment file parity, port map integrity, and lockfile consistency.

Usage:
    python worktree_validator.py
    python worktree_validator.py --json
    python worktree_validator.py --checks stale,env,ports
    python worktree_validator.py --stale-days 7
"""

import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path


SEVERITY_ERROR = "error"
SEVERITY_WARNING = "warning"
SEVERITY_INFO = "info"

ALL_CHECKS = ["stale", "missing", "branch", "env", "ports", "lockfile"]

ENV_FILES = [".env", ".env.local", ".env.development", ".env.test"]
LOCKFILES = [
    "pnpm-lock.yaml", "package-lock.json", "yarn.lock",
    "requirements.txt", "Pipfile.lock", "go.sum", "Gemfile.lock",
]


def run_git(args, cwd=None):
    """Run a git command and return stdout, or None on failure."""
    try:
        result = subprocess.run(
            ["git"] + args,
            capture_output=True, text=True, cwd=cwd, timeout=30
        )
        return result.stdout.strip() if result.returncode == 0 else None
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return None


def get_repo_root():
    """Get the main repository root."""
    root = run_git(["rev-parse", "--show-toplevel"])
    if not root:
        print("Error: not inside a git repository.", file=sys.stderr)
        sys.exit(1)
    return root


def parse_worktrees(repo_root):
    """Parse worktree list into structured records."""
    raw = run_git(["worktree", "list", "--porcelain"], cwd=repo_root)
    if not raw:
        return []

    worktrees = []
    current = {}
    for line in raw.splitlines():
        if line.startswith("worktree "):
            if current:
                worktrees.append(current)
            current = {"path": line[len("worktree "):]}
        elif line.startswith("HEAD "):
            current["head"] = line[len("HEAD "):]
        elif line.startswith("branch "):
            current["branch"] = line[len("branch "):].replace("refs/heads/", "")
        elif line == "detached":
            current["branch"] = "(detached)"
            current["detached"] = True
        elif line == "bare":
            current["bare"] = True
    if current:
        worktrees.append(current)
    return worktrees


def get_env_keys(filepath):
    """Extract variable names from an env file, ignoring comments and blanks."""
    keys = set()
    try:
        with open(filepath, "r") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key = line.split("=", 1)[0].strip()
                    if key:
                        keys.add(key)
    except (OSError, UnicodeDecodeError):
        pass
    return keys


def check_stale(worktrees, repo_root, stale_days):
    """Check for worktrees older than the stale threshold."""
    findings = []
    for wt in worktrees:
        if wt["path"] == repo_root:
            continue
        if not os.path.isdir(wt["path"]):
            continue
        try:
            mtime = os.path.getmtime(wt["path"])
            age_days = int((time.time() - mtime) / 86400)
        except OSError:
            age_days = -1

        if age_days >= stale_days:
            findings.append({
                "check": "stale",
                "severity": SEVERITY_WARNING,
                "path": wt["path"],
                "branch": wt.get("branch", "?"),
                "age_days": age_days,
                "threshold": stale_days,
                "message": f"Worktree is {age_days} days old (threshold: {stale_days})",
            })
    return findings


def check_missing(worktrees, repo_root):
    """Check for worktrees whose directories no longer exist."""
    findings = []
    for wt in worktrees:
        if wt["path"] == repo_root:
            continue
        if not os.path.isdir(wt["path"]):
            findings.append({
                "check": "missing",
                "severity": SEVERITY_ERROR,
                "path": wt["path"],
                "branch": wt.get("branch", "?"),
                "message": "Worktree directory does not exist. Run 'git worktree prune'.",
            })
    return findings


def check_branch(worktrees, repo_root):
    """Check for orphaned branches (remote deleted) and detached HEADs."""
    findings = []
    # Get list of remote branches
    remote_raw = run_git(["branch", "-r", "--format=%(refname:short)"], cwd=repo_root)
    remote_branches = set()
    if remote_raw:
        remote_branches = {b.strip() for b in remote_raw.splitlines()}

    for wt in worktrees:
        if wt["path"] == repo_root:
            continue
        branch = wt.get("branch", "")

        if wt.get("detached"):
            findings.append({
                "check": "branch",
                "severity": SEVERITY_WARNING,
                "path": wt["path"],
                "branch": branch,
                "message": "Worktree is in detached HEAD state.",
            })
            continue

        if not branch:
            continue

        # Check if the branch has a remote tracking counterpart
        tracking = run_git(
            ["rev-parse", "--abbrev-ref", f"{branch}@{{upstream}}"], cwd=repo_root
        )
        if tracking and tracking not in remote_branches:
            findings.append({
                "check": "branch",
                "severity": SEVERITY_WARNING,
                "path": wt["path"],
                "branch": branch,
                "message": f"Upstream '{tracking}' no longer exists on remote.",
            })

        # Check if the branch is merged into main
        merged_raw = run_git(["branch", "--merged", "main"], cwd=repo_root)
        if merged_raw:
            merged_list = [b.strip().lstrip("* ") for b in merged_raw.splitlines()]
            if branch in merged_list and branch != "main":
                findings.append({
                    "check": "branch",
                    "severity": SEVERITY_INFO,
                    "path": wt["path"],
                    "branch": branch,
                    "message": "Branch is already merged into main. Consider removing this worktree.",
                })

    return findings


def check_env(worktrees, repo_root):
    """Validate env file parity between main repo and worktrees."""
    findings = []

    # Collect main repo env keys per file
    main_env = {}
    for env_file in ENV_FILES:
        main_path = os.path.join(repo_root, env_file)
        if os.path.isfile(main_path):
            main_env[env_file] = get_env_keys(main_path)

    if not main_env:
        findings.append({
            "check": "env",
            "severity": SEVERITY_INFO,
            "path": repo_root,
            "message": "No .env files found in main repo. Skipping parity check.",
        })
        return findings

    for wt in worktrees:
        if wt["path"] == repo_root:
            continue
        if not os.path.isdir(wt["path"]):
            continue

        for env_file, main_keys in main_env.items():
            wt_env_path = os.path.join(wt["path"], env_file)
            if not os.path.isfile(wt_env_path):
                findings.append({
                    "check": "env",
                    "severity": SEVERITY_WARNING,
                    "path": wt["path"],
                    "file": env_file,
                    "message": f"Missing {env_file} (present in main repo).",
                })
                continue

            wt_keys = get_env_keys(wt_env_path)
            missing_in_wt = main_keys - wt_keys
            extra_in_wt = wt_keys - main_keys

            if missing_in_wt:
                findings.append({
                    "check": "env",
                    "severity": SEVERITY_WARNING,
                    "path": wt["path"],
                    "file": env_file,
                    "missing_keys": sorted(missing_in_wt),
                    "message": f"{env_file} is missing {len(missing_in_wt)} key(s) present in main: {', '.join(sorted(missing_in_wt))}",
                })
            if extra_in_wt:
                findings.append({
                    "check": "env",
                    "severity": SEVERITY_INFO,
                    "path": wt["path"],
                    "file": env_file,
                    "extra_keys": sorted(extra_in_wt),
                    "message": f"{env_file} has {len(extra_in_wt)} extra key(s) not in main: {', '.join(sorted(extra_in_wt))}",
                })

    return findings


def check_ports(worktrees, repo_root):
    """Validate .worktree-ports.json integrity and uniqueness."""
    findings = []
    all_ports = {}  # port -> worktree path

    for wt in worktrees:
        if not os.path.isdir(wt["path"]):
            continue
        port_file = os.path.join(wt["path"], ".worktree-ports.json")
        if wt["path"] == repo_root:
            if not os.path.isfile(port_file):
                continue  # main repo may not have a port file
        if not os.path.isfile(port_file):
            if wt["path"] != repo_root:
                findings.append({
                    "check": "ports",
                    "severity": SEVERITY_WARNING,
                    "path": wt["path"],
                    "message": "Missing .worktree-ports.json file.",
                })
            continue

        try:
            with open(port_file, "r") as f:
                data = json.load(f)
        except (json.JSONDecodeError, OSError) as e:
            findings.append({
                "check": "ports",
                "severity": SEVERITY_ERROR,
                "path": wt["path"],
                "message": f"Invalid .worktree-ports.json: {e}",
            })
            continue

        ports = data.get("ports", {})
        if not isinstance(ports, dict) or not ports:
            findings.append({
                "check": "ports",
                "severity": SEVERITY_ERROR,
                "path": wt["path"],
                "message": "Port map is empty or malformed in .worktree-ports.json.",
            })
            continue

        # Check for collisions with other worktrees
        for service, port in ports.items():
            if not isinstance(port, int):
                findings.append({
                    "check": "ports",
                    "severity": SEVERITY_ERROR,
                    "path": wt["path"],
                    "message": f"Port for '{service}' is not an integer: {port}",
                })
                continue
            key = f"{service}:{port}"
            if port in all_ports and all_ports[port] != wt["path"]:
                findings.append({
                    "check": "ports",
                    "severity": SEVERITY_ERROR,
                    "path": wt["path"],
                    "message": f"Port {port} ({service}) conflicts with {all_ports[port]}.",
                })
            all_ports[port] = wt["path"]

        # Validate branch matches
        recorded_branch = data.get("branch", "")
        actual_branch = wt.get("branch", "")
        if recorded_branch and actual_branch and recorded_branch != actual_branch:
            findings.append({
                "check": "ports",
                "severity": SEVERITY_WARNING,
                "path": wt["path"],
                "message": f"Branch mismatch: ports.json says '{recorded_branch}', git says '{actual_branch}'.",
            })

    return findings


def check_lockfile(worktrees, repo_root):
    """Check that lockfiles in worktrees match the main repo."""
    findings = []

    main_locks = {}
    for lf in LOCKFILES:
        main_path = os.path.join(repo_root, lf)
        if os.path.isfile(main_path):
            try:
                main_locks[lf] = os.path.getmtime(main_path)
            except OSError:
                pass

    for wt in worktrees:
        if wt["path"] == repo_root:
            continue
        if not os.path.isdir(wt["path"]):
            continue

        for lf, main_mtime in main_locks.items():
            wt_lf = os.path.join(wt["path"], lf)
            if not os.path.isfile(wt_lf):
                findings.append({
                    "check": "lockfile",
                    "severity": SEVERITY_INFO,
                    "path": wt["path"],
                    "file": lf,
                    "message": f"Lockfile {lf} not found (present in main repo).",
                })
                continue

            try:
                wt_mtime = os.path.getmtime(wt_lf)
            except OSError:
                continue

            # Lockfile in worktree is older than main
            if wt_mtime < main_mtime:
                findings.append({
                    "check": "lockfile",
                    "severity": SEVERITY_WARNING,
                    "path": wt["path"],
                    "file": lf,
                    "message": f"{lf} is older than main repo copy. Dependencies may be out of sync.",
                })

    return findings


def run_checks(checks, worktrees, repo_root, stale_days):
    """Run the requested checks and return all findings."""
    findings = []
    dispatch = {
        "stale": lambda: check_stale(worktrees, repo_root, stale_days),
        "missing": lambda: check_missing(worktrees, repo_root),
        "branch": lambda: check_branch(worktrees, repo_root),
        "env": lambda: check_env(worktrees, repo_root),
        "ports": lambda: check_ports(worktrees, repo_root),
        "lockfile": lambda: check_lockfile(worktrees, repo_root),
    }
    for check in checks:
        if check in dispatch:
            findings.extend(dispatch[check]())
    return findings


def main():
    parser = argparse.ArgumentParser(
        description="Validate worktree health: stale trees, missing branches, env parity, port conflicts."
    )
    parser.add_argument(
        "--json", action="store_true", help="Output in JSON format"
    )
    parser.add_argument(
        "--checks", type=str, default=",".join(ALL_CHECKS),
        help=f"Comma-separated checks to run (default: {','.join(ALL_CHECKS)})"
    )
    parser.add_argument(
        "--stale-days", type=int, default=14,
        help="Stale threshold in days (default: 14)"
    )

    args = parser.parse_args()
    requested = [c.strip() for c in args.checks.split(",") if c.strip()]

    invalid = [c for c in requested if c not in ALL_CHECKS]
    if invalid:
        print(f"Error: unknown check(s): {', '.join(invalid)}", file=sys.stderr)
        print(f"Available: {', '.join(ALL_CHECKS)}", file=sys.stderr)
        sys.exit(1)

    repo_root = get_repo_root()
    worktrees = parse_worktrees(repo_root)

    findings = run_checks(requested, worktrees, repo_root, args.stale_days)

    if args.json:
        output = {
            "repo_root": repo_root,
            "worktree_count": len(worktrees),
            "checks_run": requested,
            "finding_count": len(findings),
            "errors": sum(1 for f in findings if f["severity"] == SEVERITY_ERROR),
            "warnings": sum(1 for f in findings if f["severity"] == SEVERITY_WARNING),
            "info": sum(1 for f in findings if f["severity"] == SEVERITY_INFO),
            "findings": findings,
        }
        print(json.dumps(output, indent=2))
        return

    print(f"Worktree Health Report")
    print(f"Repository: {repo_root}")
    print(f"Worktrees:  {len(worktrees)}")
    print(f"Checks:     {', '.join(requested)}")
    print("=" * 70)

    if not findings:
        print("\nAll checks passed. No issues found.")
        return

    errors = [f for f in findings if f["severity"] == SEVERITY_ERROR]
    warnings = [f for f in findings if f["severity"] == SEVERITY_WARNING]
    infos = [f for f in findings if f["severity"] == SEVERITY_INFO]

    for label, group, marker in [
        ("ERRORS", errors, "X"),
        ("WARNINGS", warnings, "!"),
        ("INFO", infos, "~"),
    ]:
        if not group:
            continue
        print(f"\n  {label} ({len(group)})")
        print(f"  {'-' * 40}")
        for f in group:
            print(f"  [{marker}] {f['message']}")
            print(f"      Path: {f['path']}")
            if "branch" in f:
                print(f"      Branch: {f['branch']}")
            print()

    print(f"Summary: {len(errors)} error(s), {len(warnings)} warning(s), {len(infos)} info")
    if errors:
        sys.exit(1)


if __name__ == "__main__":
    main()