#!/usr/bin/env python3
"""
Aggressive cleanup of stale Claude Code background-session worktrees.

Unlike safe `cleanup-merged-worktrees.sh`, this queries Forgejo PR API to
detect squash-merged branches (where local SHA isn't ancestor of main
because Forgejo squash created a new commit).

Decision matrix per worktree branch:

    PR state              | Action
    ----------------------|---------------------------------------------
    merged (any style)    | REMOVE
    closed (no merge)     | REMOVE (work abandoned)
    open                  | KEEP (active PR in progress)
    no PR + age <14d      | KEEP (recent WIP, no PR yet)
    no PR + age >=14d     | REMOVE (stale orphan)

Requires:
    - FORGEJO_TOKEN env var (or in ~/.claude/settings.json env block)
    - FORGEJO_URL env var (default: https://git.gendsgn.ru)

Usage:
    python scripts/cleanup-worktrees-aggressive.py --dry-run     # show plan
    python scripts/cleanup-worktrees-aggressive.py               # execute
    AGE_DAYS=30 python scripts/cleanup-worktrees-aggressive.py   # override threshold
"""

from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
import time
import urllib.error
import urllib.request
from collections import Counter
from pathlib import Path

OWNER = "lekss361"
REPO_NAME = "gendesign"
FORGEJO_URL = os.environ.get("FORGEJO_URL", "https://git.gendsgn.ru")
TOKEN = os.environ.get("FORGEJO_TOKEN")
AGE_DAYS = int(os.environ.get("AGE_DAYS", "14"))
DRY_RUN = "--dry-run" in sys.argv

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


def run(cmd: list[str], cwd: Path | None = None, check: bool = True) -> str:
    """Run a subprocess and return stdout, swallowing errors when check=False."""
    result = subprocess.run(
        cmd,
        cwd=cwd or REPO_ROOT,
        capture_output=True,
        text=True,
        check=False,
    )
    if check and result.returncode != 0:
        raise RuntimeError(f"{' '.join(cmd)}\n{result.stderr}")
    return result.stdout


def fetch_all_prs() -> dict[str, dict]:
    """Returns {branch_name: {state, merged}} for all PRs across all states."""
    if not TOKEN:
        sys.exit(
            "ERROR: FORGEJO_TOKEN env var not set.\n"
            "Hint: $env:FORGEJO_TOKEN = '<token>' (PowerShell) or export it before running."
        )

    prs: dict[str, dict] = {}
    page = 1
    api = f"{FORGEJO_URL}/api/v1/repos/{OWNER}/{REPO_NAME}/pulls"

    while page <= 30:
        url = f"{api}?state=all&limit=50&page={page}&sort=newest"
        req = urllib.request.Request(
            url, headers={"Authorization": f"token {TOKEN}"}
        )
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                batch = json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as e:
            sys.exit(f"ERROR: Forgejo API {url} → {e.code} {e.reason}")
        except Exception as e:
            sys.exit(f"ERROR: fetch {url} failed: {e}")

        if not batch:
            break

        for pr in batch:
            branch = pr.get("head", {}).get("ref")
            if not branch:
                continue
            # Newest-wins: only set if not already seen (sort=newest gives
            # most recent first, so first occurrence is the newest PR for
            # that branch)
            prs.setdefault(
                branch,
                {
                    "state": pr.get("state"),
                    "merged": bool(pr.get("merged")),
                    "number": pr.get("number"),
                },
            )

        page += 1

    return prs


def get_worktrees() -> list[tuple[str, str]]:
    """Returns list of (path, branch_short_name) tuples, skipping main."""
    out = run(["git", "worktree", "list", "--porcelain"])
    worktrees: list[tuple[str, str]] = []
    current = {}
    for line in out.splitlines():
        if line.startswith("worktree "):
            current = {"path": line[len("worktree ") :]}
        elif line.startswith("branch "):
            current["branch"] = line[len("branch ") :]
        elif line == "" and current:
            branch = current.get("branch", "")
            if branch.startswith("refs/heads/"):
                short = branch[len("refs/heads/") :]
                worktrees.append((current["path"], short))
            current = {}
    if current and "branch" in current:
        branch = current["branch"]
        if branch.startswith("refs/heads/"):
            worktrees.append((current["path"], branch[len("refs/heads/") :]))
    return worktrees


def branch_age_days(path: str) -> int | None:
    """Returns age in days of branch HEAD, or None on error."""
    try:
        ts = run(["git", "-C", path, "log", "-1", "--format=%ct"], check=False).strip()
        if not ts:
            return None
        return (int(time.time()) - int(ts)) // 86400
    except Exception:
        return None


def main() -> int:
    if DRY_RUN:
        print("DRY RUN — no worktrees will be removed\n")

    print("Fetching all PRs from Forgejo (paginated, up to 1500)...")
    prs = fetch_all_prs()
    print(f"Indexed {len(prs)} unique branch→PR mappings.\n")

    worktrees = get_worktrees()
    print(f"Scanning {len(worktrees)} worktrees...\n")

    categories: Counter[str] = Counter()
    to_remove: list[tuple[str, str, str]] = []  # (path, branch, reason)
    repo_root_str = str(REPO_ROOT)

    for path, branch in worktrees:
        # Don't touch the worktree we're standing in.
        if os.path.normcase(os.path.normpath(path)) == os.path.normcase(
            os.path.normpath(repo_root_str)
        ):
            categories["skipped"] += 1
            continue

        pr = prs.get(branch)
        if pr:
            if pr["merged"]:
                cat = "merged"
                reason = f"remove (merged PR #{pr['number']})"
            elif pr["state"] == "closed":
                cat = "closed"
                reason = f"remove (closed PR #{pr['number']}, no merge)"
            else:
                cat = "open"
                reason = f"keep (open PR #{pr['number']})"
        else:
            age = branch_age_days(path)
            if age is None:
                cat = "other"
                reason = "keep (no PR, age unknown)"
            elif age >= AGE_DAYS:
                cat = "stale"
                reason = f"remove (no PR, stale {age}d ≥ {AGE_DAYS}d)"
            else:
                cat = "recent"
                reason = f"keep (no PR, recent {age}d)"

        print(f"  [{cat:>6}] {branch} → {reason}")
        categories[cat] += 1

        if cat in {"merged", "closed", "stale"}:
            to_remove.append((path, branch, reason))

    if not DRY_RUN:
        print(f"\nRemoving {len(to_remove)} worktree(s)...")
        failed = 0
        for path, branch, _ in to_remove:
            # `--force --force` (double flag) overrides Claude Code supervisor
            # locks (`claude agent <id> (pid X)` lock messages).
            result = subprocess.run(
                ["git", "worktree", "remove", "--force", "--force", path],
                capture_output=True,
                text=True,
                cwd=REPO_ROOT,
            )
            if result.returncode != 0:
                failed += 1
                err = result.stderr.strip().splitlines()[-1] if result.stderr else "unknown"
                print(f"  FAILED: {branch} → {err}", file=sys.stderr)
        run(["git", "worktree", "prune"], check=False)
        if failed:
            print(f"\nWARN: {failed} worktree(s) failed to remove (see stderr above)")

        # Phase 2: orphan-dir cleanup. After `git worktree remove`, git no longer
        # tracks the dir, but filesystem may still have it (Windows file locks
        # held by Claude Code supervisor, etc). Try to remove via shutil.
        wt_root = REPO_ROOT / ".claude" / "worktrees"
        if wt_root.is_dir():
            fs_dirs = {p.name for p in wt_root.iterdir() if p.is_dir()}
            git_dirs = {Path(p).name for p, _ in get_worktrees()}
            orphans = fs_dirs - git_dirs
            if orphans:
                print(f"\nFound {len(orphans)} orphan dir(s) (git-untracked, FS leftover)...")
                orphan_removed = 0
                orphan_locked = 0
                for o in orphans:
                    try:
                        shutil.rmtree(wt_root / o, ignore_errors=False)
                        orphan_removed += 1
                    except PermissionError:
                        orphan_locked += 1
                    except Exception as e:
                        print(f"  ORPHAN FAILED: {o} → {type(e).__name__}", file=sys.stderr)
                print(f"  Removed {orphan_removed}, locked (need Claude Code restart) {orphan_locked}")
                if orphan_locked:
                    print(
                        f"\n  Hint: {orphan_locked} dirs held by Claude Code supervisor "
                        f"(open file handles). After next Claude Code restart, re-run this "
                        f"script — they will be removable then."
                    )

    print("\nSummary:")
    print(f"  Remove — merged:   {categories['merged']}")
    print(f"  Remove — closed:   {categories['closed']}")
    print(f"  Remove — stale:    {categories['stale']}  (no PR, age >={AGE_DAYS}d)")
    print(f"  Keep   — open PR:  {categories['open']}")
    print(
        f"  Keep   — recent:   {categories['recent']}  (no PR yet, age <{AGE_DAYS}d)"
    )
    print(f"  Keep   — other:    {categories['other']}  (no PR, unknown age)")
    print(f"  Skipped (current): {categories['skipped']}")

    total_remove = categories["merged"] + categories["closed"] + categories["stale"]
    print()
    if DRY_RUN:
        print(
            f"DRY RUN — would remove {total_remove} worktree(s). "
            "Re-run without --dry-run to execute."
        )
    else:
        after = len(get_worktrees()) + 1  # +1 for main
        print(f"Worktrees remaining: {after}")

    return 0


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