#!/usr/bin/env python3
"""polytree — one feature, many repos.

Create or link a git worktree with the SAME branch across several repositories,
then launch ONE agent session that can see and edit all of them at once.

Tool-agnostic by design:
  * worktrees : plain `git worktree`. Orca is an OPTIONAL backend (auto-detected).
  * agent     : an adapter table in config. `claude` and `codex` ship built in;
                any other agent that accepts extra roots works via config only.

Config: $XDG_CONFIG_HOME/polytree/config.toml  (see `polytree --help`)
License: MIT
"""
from __future__ import annotations

import argparse
import json
import os
import re
import shlex
import shutil
import signal
import subprocess
import sys
import tomllib
from pathlib import Path

VERSION = "0.1.8"
VERBOSE = False  # set by --verbose; when on, every git/orca command is echoed
CONFIG = Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config")) / "polytree" / "config.toml"

# Built-in agent adapters. Anything here can be overridden/extended from config.
#
# Verified behaviour (Claude Code 2.1.209) for directories attached with --add-dir:
#   .claude/skills/  -> discovered automatically
#   CLAUDE.md        -> NOT loaded unless CLAUDE_CODE_ADDITIONAL_DIRECTORIES_CLAUDE_MD=1
#   .mcp.json        -> NOT loaded; --mcp-config pulls it in (and adds to existing)
#   settings/hooks   -> never loaded from an attached dir (host repo only)
# Codex reads AGENTS.md hierarchically across roots and needs no env var.
BUILTIN_AGENTS: dict[str, dict] = {
    "claude": {
        "attach": "--add-dir {dir}",
        "env": {"CLAUDE_CODE_ADDITIONAL_DIRECTORIES_CLAUDE_MD": "1"},
        "attach_if_exists": {".mcp.json": "--mcp-config {dir}/.mcp.json"},
    },
    "codex": {
        "attach": "--add-dir {dir}",
    },
}

EXAMPLE = """\
example config (~/.config/polytree/config.toml):

  backend = "auto"          # auto | git | orca  (auto = orca if installed, else git)
  agent   = "claude"        # any key under [agents], or a built-in (claude, codex)
  root    = "~/polytree"    # git backend: worktrees live at <root>/<feature>/<repo>
  base    = "origin/main"   # optional global default base ref

  [[repos]]                 # the first repo (or host = true) hosts the agent
  path = "~/code/my-api"
  base = "origin/dev"       # optional; overrides the global default
  host = true

  [[repos]]
  path = "~/code/my-web"

  [agents.my-agent]         # add any agent without touching the code
  cmd    = "my-agent --fast"                 # optional; defaults to the table key
  attach = "--root {dir}"                    # required: how it takes an extra dir
  env    = { MY_AGENT_EXTRA = "1" }          # optional
  attach_if_exists = { ".mcp.json" = "--mcp-config {dir}/.mcp.json" }   # optional
"""

ENV_KEY = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z")


class Fail(Exception):
    """A user-facing error. Raised (not exited) so callers can roll back first."""


def die(msg: str):
    raise Fail(msg)


def info(msg: str) -> None:
    print(msg, flush=True)


def confirm(prompt: str) -> bool:
    """Ask a yes/no question. Refuses (returns False) when there is no terminal
    to ask — so a non-interactive run never blocks or assumes 'yes'."""
    if not sys.stdin.isatty():
        return False
    try:
        return input(f"{prompt} [y/N] ").strip().lower() in ("y", "yes")
    except EOFError:
        return False


# ----------------------------------------------------------------------- process
def trace(args: list[str], cwd: str | None = None) -> None:
    if VERBOSE:
        line = " ".join(shlex.quote(a) for a in args)
        print(f"+ {line}" + (f"   (cwd={cwd})" if cwd else ""), file=sys.stderr, flush=True)


def run(args: list[str], cwd: str | None = None, check: bool = True) -> str:
    trace(args, cwd)
    p = subprocess.run(args, cwd=cwd, capture_output=True, text=True)
    if check and p.returncode != 0:
        die(f"command failed: {' '.join(args)}\n{p.stderr.strip()}")
    return p.stdout


def orca_error(stdout: str) -> str | None:
    """The message Orca reports in its JSON, or None if it did not report one.

    Orca prints failures to stdout as {"ok": false, "error": {"message": ...}}.
    The generic runner only ever shows stderr, which Orca leaves empty — so on
    its own a failed Orca call looks like a bare command with no reason.
    """
    try:
        payload = json.loads(stdout)
    except (ValueError, TypeError):
        return None
    if isinstance(payload, dict) and not payload.get("ok", True):
        return (payload.get("error") or {}).get("message") or "Orca reported a failure"
    return None


def orca_hint(reason: str) -> str:
    if reason == "repo_not_found" or any(
        s in reason.lower() for s in ("not found", "not registered", "unknown repo", "no such repo")
    ):
        return ('\n  That repo is not registered in Orca. Add it with '
                '`orca repo add --path <path>`, or set backend = "git" to skip Orca.')
    return ""


def orca_try(args: list[str], cwd: str | None = None) -> tuple[str, str | None]:
    """Run an Orca command; return (stdout, reason). `reason` is why it failed, or
    None. Never dies — the caller decides. Orca reports failures as JSON on stdout
    ({"ok": false, "error": {...}}); the generic runner only sees stderr (empty)."""
    trace(args, cwd)
    p = subprocess.run(args, cwd=cwd, capture_output=True, text=True)
    reason = orca_error(p.stdout)
    if p.returncode != 0 and not reason:
        reason = p.stderr.strip() or "command failed"
    return p.stdout, reason


def orca_run(args: list[str], cwd: str | None = None) -> str:
    """Run an Orca command, surfacing the reason it gives when it fails."""
    out, reason = orca_try(args, cwd)
    if reason:
        die(f"orca: {reason}{orca_hint(reason)}")
    return out


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


def git_ok(repo: str, *args: str) -> bool:
    trace(["git", "-C", repo, *args])
    return subprocess.run(["git", "-C", repo, *args], capture_output=True, text=True).returncode == 0


def orca_bin() -> str | None:
    return shutil.which("orca-ide") or shutil.which("orca")


# ------------------------------------------------------------------------ config
def load_config() -> dict:
    if not CONFIG.exists():
        die(f"no config at {CONFIG}\n\n{EXAMPLE}")
    try:
        with CONFIG.open("rb") as fh:
            cfg = tomllib.load(fh)
    except tomllib.TOMLDecodeError as e:
        die(f"invalid TOML in {CONFIG}: {e}")

    repos = cfg.get("repos") or []
    if len(repos) < 2:
        die(f"need at least 2 [[repos]] in {CONFIG}\n\n{EXAMPLE}")

    seen: dict[str, str] = {}
    paths: set[str] = set()
    for r in repos:
        if "path" not in r:
            die("every [[repos]] entry needs a `path`")
        p = Path(r["path"]).expanduser()
        if not p.is_dir():
            die(f"repo path does not exist: {p}")
        r["path"] = str(p.resolve())
        if not git_ok(r["path"], "rev-parse", "--git-dir"):
            die(f"not a git repo: {r['path']}")
        if r["path"] in paths:
            die(f"duplicate repo path: {r['path']}")
        paths.add(r["path"])

        r["name"] = r.get("name") or p.name
        # `name` is a directory component under <root>/<branch>/ — keep it one.
        if "/" in r["name"] or r["name"] in ("", ".", ".."):
            die(f"invalid repo name {r['name']!r}: must be a plain directory name")
        if r["name"] in seen:
            die(
                f"duplicate repo name {r['name']!r} ({seen[r['name']]} and {r['path']}).\n"
                f'Set an explicit `name = "..."` on one of them.'
            )
        seen[r["name"]] = r["path"]

    if sum(1 for r in repos if r.get("host")) > 1:
        die("only one [[repos]] entry may set `host = true`")
    repos.sort(key=lambda r: not r.get("host", False))  # host first
    cfg["repos"] = repos

    agents = {k: dict(v) for k, v in BUILTIN_AGENTS.items()}
    for name, spec in (cfg.get("agents") or {}).items():
        agents[name] = {**agents.get(name, {}), **spec}
    cfg["agents"] = agents

    backend = cfg.get("backend", "auto")
    # Case-insensitive: "Git"/"ORCA" are the same intent, and a rejected config
    # over a stray capital is a poor welcome. The value is a fixed keyword, not
    # a name, so normalising is safe.
    if isinstance(backend, str):
        backend = backend.strip().lower()
    if backend not in ("auto", "git", "orca"):
        die(f"invalid backend {backend!r} (auto | git | orca)")
    if backend == "auto":
        backend = "orca" if orca_bin() else "git"
    elif backend == "orca" and not orca_bin():
        die("backend 'orca' requested but the Orca CLI is not installed")
    cfg["backend"] = backend
    cfg["agent"] = cfg.get("agent", "claude")
    cfg["root"] = str(Path(cfg.get("root", "~/polytree")).expanduser())
    cfg["subshell"] = bool(cfg.get("subshell", True))  # git backend: land in a shell in the worktree (default on)
    return cfg


def resolve_agent(cfg: dict) -> dict:
    """Validate the agent adapter. Call BEFORE creating anything."""
    name = cfg["agent"]
    spec = cfg["agents"].get(name)
    if spec is None:
        known = ", ".join(sorted(cfg["agents"]))
        die(f"unknown agent {name!r}. Define [agents.{name}] in the config, or use one of: {known}")
    if not spec.get("attach"):
        die(f'[agents.{name}] needs `attach`, e.g. attach = "--add-dir {{dir}}"')
    argv = shlex.split(spec.get("cmd", name))
    if not argv:
        die(f"[agents.{name}] has an empty `cmd`")
    if not shutil.which(argv[0]):
        die(f"agent binary {argv[0]!r} not found in PATH")
    for k in spec.get("env") or {}:
        if not ENV_KEY.match(k):
            die(f"[agents.{name}] invalid env var name {k!r}")
    return spec


# --------------------------------------------------------------------- discovery
def worktree_records(repo: str) -> list[dict[str, str]]:
    """Every worktree of `repo` as its raw porcelain attributes.

    -z so paths containing newlines survive.
    """
    records = []
    for record in git(repo, "worktree", "list", "--porcelain", "-z").split("\0\0"):
        attrs: dict[str, str] = {}
        for field in record.split("\0"):
            if field:
                key, _, val = field.partition(" ")
                attrs[key] = val
        if attrs.get("worktree"):
            records.append(attrs)
    return records


def worktrees_of(repo: str) -> dict[str, str]:
    """branch -> worktree path, for every usable worktree of `repo`.

    Skips bare/detached worktrees (no branch) and prunable ones (already gone).
    """
    return {
        r["branch"].removeprefix("refs/heads/"): r["worktree"]
        for r in worktree_records(repo)
        if r.get("branch") and "prunable" not in r
    }


def is_locked(repo: str, path: str) -> bool:
    return any("locked" in r for r in worktree_records(repo) if r["worktree"] == path)


def branch_exists(repo: dict, branch: str) -> bool:
    return git_ok(repo["path"], "show-ref", "--verify", "--quiet", f"refs/heads/{branch}")


def sibling_map(cfg: dict, branch: str) -> list[tuple[dict, str | None]]:
    """(repo, worktree path or None) for every configured repo, host first."""
    return [(r, worktrees_of(r["path"]).get(branch)) for r in cfg["repos"]]


def current_branch() -> str:
    p = subprocess.run(["git", "rev-parse", "--abbrev-ref", "HEAD"], capture_output=True, text=True)
    if p.returncode != 0:
        die("not inside a git repo; pass the branch explicitly")
    branch = p.stdout.strip()
    if branch == "HEAD":
        die("you are on a detached HEAD; pass the branch explicitly")
    return branch


def check_branch_name(branch: str) -> None:
    if branch.startswith("-") or not git_ok(".", "check-ref-format", "--branch", branch):
        die(f"invalid branch name: {branch!r}")


def pick_host(pairs: list[tuple[dict, str]], host_opt: str | None) -> tuple[dict, str]:
    """Host = the first configured repo, unless --host says otherwise.

    Deliberately NOT the cwd: hooks and settings only ever load from the host, so
    the host must be predictable rather than depend on where you happen to stand.
    """
    if host_opt:
        for r, p in pairs:
            if r["name"] == host_opt:
                return r, p
        names = ", ".join(r["name"] for r, _ in pairs)
        die(f"--host {host_opt!r} is not a repo with a worktree for this branch (have: {names})")
    return pairs[0]


# ------------------------------------------------------------------------ create
def dest_for(cfg: dict, repo: dict, branch: str) -> str:
    return str(Path(cfg["root"]) / branch / repo["name"])


def branch_of_pr(cfg: dict, spec: str) -> tuple[dict, str, str]:
    """Resolve `[<repo>#]<number>` to (repo, number, head branch) via the GitHub CLI.

    Working on a PR needs nothing from Orca: a PR is a branch, and `gh` maps one
    to the other. (Orca's own PR badge is metadata in its store, which only its
    UI can set — git has no such concept, so polytree does not pretend to.)

    The repo qualifier is not decoration: PR numbers are per-repo, so the same
    number is a different PR in each one, and guessing would open the wrong work.
    """
    if not shutil.which("gh"):
        die("--pr needs the GitHub CLI (`gh`) on PATH")
    repo_name, _, number = spec.rpartition("#")
    if not number.isdigit():
        die(f"--pr: expected <number> or <repo>#<number>, got {spec!r}")
    if repo_name:
        repo = next((r for r in cfg["repos"] if r["name"] == repo_name), None)
        if not repo:
            known = ", ".join(r["name"] for r in cfg["repos"])
            die(f"--pr: unknown repo {repo_name!r} (configured: {known})")
    else:
        repo = cfg["repos"][0]  # the host; always named in the output below
    out = run(["gh", "pr", "view", number, "--json", "headRefName", "-q", ".headRefName"],
              cwd=repo["path"], check=False).strip()
    if not out:
        die(f"could not resolve PR #{number} in {repo['name']} — is `gh` authenticated for it?")
    return repo, number, out


def configured_base(cfg: dict, repo: dict, override: str | None) -> str | None:
    """The base someone asked for: --base, then the repo's, then the global one.

    None means nobody asked — the git backend then falls back to the repo's
    origin/HEAD, and the orca backend defers to Orca's own base ref.
    """
    return override or repo.get("base") or cfg.get("base")


def default_base(repo: dict) -> str:
    head = git(repo["path"], "symbolic-ref", "--quiet", "refs/remotes/origin/HEAD", check=False).strip()
    if head:
        return head.removeprefix("refs/remotes/")
    local = git(repo["path"], "symbolic-ref", "--quiet", "--short", "HEAD", check=False).strip()
    if local:
        return local
    die(
        f"cannot determine a base ref for {repo['name']}: set `base` in the config, "
        f"or run `git -C {repo['path']} remote set-head origin -a`"
    )


def create_git(cfg: dict, repo: dict, branch: str, override: str | None = None, existing: bool = False) -> str:
    dest = dest_for(cfg, repo, branch)
    if existing:  # check the branch out as-is; it has whatever upstream it has
        git(repo["path"], "worktree", "add", dest, branch)
        return dest
    base = configured_base(cfg, repo, override) or default_base(repo)
    # --no-track matters: branching off a remote ref like origin/dev would, with
    # git's default branch.autoSetupMerge, make the feature branch track dev. Then
    # `git push` fails and helpfully suggests `git push origin HEAD:dev` — pushing
    # unreviewed work onto the shared branch. Orca's backend sets no upstream
    # either, so this also keeps the two backends consistent.
    git(repo["path"], "worktree", "add", "-b", branch, "--no-track", dest, base)
    return dest


def create_orca(cfg: dict, repo: dict, branch: str, override: str | None = None,
                existing: bool = False, issue: str | None = None) -> str:
    """Create through Orca, then locate the worktree with git.

    We deliberately do NOT parse Orca's JSON for the path: that schema is not a
    stable contract, and guessing wrong would silently point the agent at the
    main checkout. Orca worktrees are ordinary git worktrees, so git is truth.
    """
    sel = repo.get("orca") or f"path:{repo['path']}"
    # --no-parent: Orca otherwise makes the new worktree a child of whatever
    # worktree you happen to be standing in, so the same command produces
    # different lineage depending on your cwd — and only for the repo that
    # matches it, leaving the set inconsistent with itself. A feature set is
    # independent work, not a branch stacked on the caller's. (Lineage only;
    # the git base is unaffected.)
    args = [orca_bin(), "worktree", "create", "--repo", sel, "--name", branch,
            "--no-parent", "--json"]
    if issue:
        args += ["--issue", str(issue)]
    if not existing:
        base = configured_base(cfg, repo, override)
        if base:  # nobody asked -> let Orca use the base ref configured for the repo
            args += ["--base-branch", base]

    # Find the new worktree by diffing, not by looking up `branch`: Orca slugifies
    # --name into the branch it creates ("feature/x" -> "feature-x"), so the branch
    # we asked for may not be the branch we get. Snapshot branches too — we have to
    # know whether Orca invented the branch or grabbed one that was already there.
    before = {r["worktree"] for r in worktree_records(repo["path"])}
    branches_before = local_branches(repo)
    out, reason = orca_try(args)
    if reason == "repo_not_found" and sel.startswith("path:"):
        # Orca doesn't know this repo yet. Offer to register it and retry, rather
        # than dying and leaving the user to run `orca repo add` by hand.
        if confirm(f"{repo['name']} is not registered in Orca. Add it (orca repo add {repo['path']})?"):
            orca_run([orca_bin(), "repo", "add", "--path", repo["path"], "--json"])
            out, reason = orca_try(args)
    if reason:
        die(f"orca: {reason}{orca_hint(reason)}")
    fresh = [r for r in worktree_records(repo["path"]) if r["worktree"] not in before]
    if len(fresh) != 1:
        die(
            f"expected Orca to create exactly one worktree in {repo['name']}, saw {len(fresh)}.\n"
            f"Check `orca worktree list`; polytree cannot tell which one is new."
        )
    path, made = fresh[0]["worktree"], fresh[0].get("branch", "").removeprefix("refs/heads/")
    ours = bool(made) and made not in branches_before  # did Orca invent this branch?

    # Everything below can fail (a rename can hit a refs D/F conflict; a checkout
    # can be ambiguous). Orca's worktree already exists by now, and the caller's
    # rollback cannot see it yet — so clean up here or it leaks and the "rolled
    # back N" line becomes a lie.
    try:
        if existing:
            # Orca cannot check an existing branch out: it always makes a new one
            # off the base, which would silently hand you an empty branch instead
            # of the work you meant to review. Put it on the right branch.
            #
            # Drop Orca's branch BEFORE checking out, and detach first so git will
            # let go of it. Order matters: when the branch we want has no slug to
            # change (no slash) and only exists on the remote, Orca's throwaway is
            # named exactly like it — so checking out first is a no-op that leaves
            # the worktree sitting on base, and the delete then fails because the
            # branch is checked out. That is silent, and it is the common shape of
            # a PR branch.
            if ours:
                git(path, "checkout", "--detach", "--quiet")
                git(repo["path"], "branch", "-D", made)  # loud: we made this one
            git(path, "checkout", branch)  # DWIMs a local branch from origin/<branch>
        elif made and made != branch:
            if not ours:
                # Orca slugified our name onto a branch that already existed and
                # reused it. Renaming would quietly steal someone's branch.
                die(
                    f"Orca turns the name {branch!r} into the branch {made!r}, which already "
                    f"exists in {repo['name']}, and reused it instead of making a new one.\n"
                    f"Rename or delete {made!r}, or pick another name."
                )
            git(repo["path"], "branch", "-m", made, branch)  # undo the slugification
    except BaseException:
        if ours:
            remove_worktree(cfg, repo, path, force=True)
            git(repo["path"], "branch", "-D", made, check=False)
        else:
            remove_worktree_keeping_branch(cfg, repo, path, made)
        raise
    return path


def is_dirty(worktree: str) -> bool:
    p = subprocess.run(["git", "-C", worktree, "status", "--porcelain"], capture_output=True, text=True)
    if p.returncode != 0:
        return True  # can't tell -> assume there is something to lose
    return bool(p.stdout.strip())


def has_populated_submodules(worktree: str) -> bool:
    """git refuses to remove a worktree with populated submodules, even a clean
    one — and `status --porcelain` says nothing about it."""
    p = subprocess.run(["git", "-C", worktree, "submodule", "status"], capture_output=True, text=True)
    if p.returncode != 0:
        return False
    # "-<sha>" means registered but not populated; anything else is populated.
    return any(line and not line.startswith("-") for line in p.stdout.splitlines())


def removal_blocker(repo: dict, path: str, force: bool) -> str | None:
    """Why `path` must not be removed, or None.

    Covers the conditions git refuses on: the main checkout, locked worktrees,
    uncommitted changes, and populated submodules. This list is an approximation
    of git's rules, not a contract — so cmd_rm also reports honestly if a removal
    fails anyway, rather than trusting this to be exhaustive.
    """
    if path == repo["path"]:
        return "is the main checkout"  # git refuses this even with --force
    if force:
        return None  # `remove --force --force` handles locked, dirty and submodules
    if is_locked(repo["path"], path):
        return "is locked"
    if is_dirty(path):
        return "has uncommitted changes"
    if has_populated_submodules(path):
        return "contains populated submodules"
    return None


def branch_sha(repo: dict, branch: str) -> str:
    return git(repo["path"], "rev-parse", "--verify", f"refs/heads/{branch}", check=False).strip()


def remove_worktree_keeping_branch(cfg: dict, repo: dict, path: str, branch: str) -> None:
    """Remove a worktree without letting the backend take the branch with it.

    `orca worktree rm` deletes the branch too, which is fine for one we just
    created and catastrophic for one that was already there.
    """
    sha = branch_sha(repo, branch)
    try:
        remove_worktree(cfg, repo, path, force=True)
    finally:
        # finally, not after: if the removal dies once Orca has already dropped
        # the branch, the restore must still happen or the branch is gone.
        if sha and not branch_exists(repo, branch):
            git(repo["path"], "branch", branch, sha, check=False)


def local_branches(repo: dict) -> set[str]:
    out = git(repo["path"], "for-each-ref", "--format=%(refname:short)", "refs/heads", check=False)
    return {b for b in out.splitlines() if b}


def remove_worktree(cfg: dict, repo: dict, path: str, force: bool) -> None:
    """Remove one worktree, and say so honestly if it did not work."""
    if cfg["backend"] == "orca":
        # Only force Orca when the user asked: don't bypass its own safeguards.
        flags = ["--force"] if force else []
        run([orca_bin(), "worktree", "rm", "--worktree", f"path:{path}", *flags, "--json"], check=False)
    if Path(path).exists():  # Orca may have removed it already
        # Doubled --force is what git requires to remove a locked worktree.
        flags = ["--force", "--force"] if force else []
        p = subprocess.run(
            ["git", "-C", repo["path"], "worktree", "remove", *flags, path], capture_output=True, text=True
        )
        if p.returncode != 0 and Path(path).exists():
            die(f"{repo['name']}: could not remove {path}\n{p.stderr.strip()}")
    git(repo["path"], "worktree", "prune", check=False)
    prune_empty_dirs(cfg, path)


def prune_empty_dirs(cfg: dict, path: str) -> None:
    """Drop the now-empty <root>/<branch>/ shell we would otherwise leave behind."""
    try:
        root = Path(cfg["root"]).resolve()
        parent = Path(path).resolve().parent
        while root in parent.parents and parent.is_dir() and not any(parent.iterdir()):
            parent.rmdir()
            parent = parent.parent
    except OSError:
        pass


def rollback(cfg: dict, created: list[tuple[dict, str]], branch: str, ours: set[str]) -> None:
    """Undo what we made. Best effort: never mask the original error.

    `ours` names the repos where polytree created the branch. Everywhere else the
    branch was already there — someone's work, `--existing`/`--pr` — and rolling
    back must take the worktree without taking their branch with it.
    """
    for repo, path in reversed(created):
        try:
            if repo["name"] in ours:
                remove_worktree(cfg, repo, path, force=True)  # safe: made seconds ago
                git(repo["path"], "branch", "-D", branch, check=False)
            else:
                remove_worktree_keeping_branch(cfg, repo, path, branch)
        except Fail:
            pass


# ------------------------------------------------------------------------ launch
def expand(template: str, directory: str) -> list[str]:
    # replace(), not format(): a template may legitimately contain braces.
    return [tok.replace("{dir}", directory) for tok in shlex.split(template)]


def build_argv(cfg: dict, spec: dict, others: list[str], prompt: str | None = None) -> tuple[list[str], dict[str, str]]:
    argv = shlex.split(spec.get("cmd", cfg["agent"]))
    for d in others:
        argv += expand(spec["attach"], d)
        for filename, template in (spec.get("attach_if_exists") or {}).items():
            if (Path(d) / filename).exists():
                argv += expand(template, d)
    if prompt:  # claude/codex take an initial prompt as a positional argument
        argv.append(prompt)
    return argv, {k: str(v) for k, v in (spec.get("env") or {}).items()}


def record_dir(path: str) -> None:
    """Leave the launch directory where a shell wrapper can find it.

    A process cannot change its parent shell's cwd, so on the git backend the
    shell lands back where you launched from once the agent exits — not in the
    worktree. The optional `pt()` shell function (see the README) passes
    POLYTREE_CD_FILE and cds into whatever we write here.
    """
    cd_file = os.environ.get("POLYTREE_CD_FILE")
    if cd_file:
        try:
            Path(cd_file).write_text(path)
        except OSError:
            pass


def launch(cfg: dict, spec: dict, host: str, others: list[str], prompt: str | None = None) -> None:
    argv, env = build_argv(cfg, spec, others, prompt)
    info(f"→ {cfg['agent']} in:")
    info(f"    host : {host}")
    for d in others:
        info(f"    +dir : {d}")
    record_dir(host)

    if cfg["backend"] == "orca":
        # Orca runs this string through a shell, hence the quoting and `K=v cmd`.
        cmd = " ".join([f"{k}={shlex.quote(v)}" for k, v in env.items()] + [shlex.quote(a) for a in argv])
        orca_run([orca_bin(), "terminal", "create", "--worktree", f"path:{host}", "--command", cmd, "--json"])
        info("    (Orca-managed terminal)")
    else:
        os.chdir(host)
        full_env = {**os.environ, **env}
        # shell=true: run the agent, then hand control to a shell in the worktree,
        # so you land there instead of back where you launched (a process can't
        # move its parent shell; a subshell is the portable way — any terminal, no
        # shell-rc edits). But NOT if we're already in a polytree shell: relaunching
        # (e.g. `link` after the agent died) then just runs the agent and returns to
        # that same shell, so shells never stack no matter how often you re-launch.
        if cfg.get("subshell") and sys.stdin.isatty() and not os.environ.get("POLYTREE_SHELL"):
            # Ignore SIGINT while the agent runs, the way a shell does — otherwise a
            # Ctrl-C reaching us (a cooked-mode agent) would abort here, kill the
            # agent and skip the shell. A *caught* handler (not SIG_IGN, which would
            # survive exec and leave the agent unable to be interrupted) resets to
            # SIG_DFL in the child on exec, so the agent still handles Ctrl-C itself.
            prev = signal.signal(signal.SIGINT, lambda *_: None)
            try:
                subprocess.run(argv, env=full_env)
            finally:
                signal.signal(signal.SIGINT, prev)  # restore before the shell inherits it
            info(f"\n→ shell in {host}  (exit to return)")
            sh = os.environ.get("SHELL") or "/bin/sh"
            if shutil.which(sh) is None:  # $SHELL set but bogus -> don't traceback
                sh = "/bin/sh"
            os.execvpe(sh, [sh], {**os.environ, "POLYTREE_SHELL": "1"})
        os.execvpe(argv[0], argv, full_env)


# ---------------------------------------------------------------------- commands
def cmd_new(cfg: dict, a: argparse.Namespace) -> None:
    pr_repo = None
    if a.pr:
        if a.name:
            die("--pr takes the branch from the pull request; drop the name argument")
        pr_repo, number, branch = branch_of_pr(cfg, a.pr)
        a.existing = True
        # Name the repo every time: PR numbers collide across repos, and a silent
        # default would cheerfully open a different PR than the one you meant.
        info(f"PR #{number} in {pr_repo['name']} → branch {branch!r}")
    elif a.name:
        branch = a.name
    else:
        die("give a branch name, or --pr <number> to take it from a pull request")

    check_branch_name(branch)
    if a.existing and a.base:
        flag = "--pr" if a.pr else "--existing"
        die(f"{flag} checks out a branch that already exists; --base has nothing to apply to")
    if a.issue and cfg["backend"] != "orca":
        die("--issue links the worktree to an issue in Orca; it needs the orca backend")
    if a.agent:
        cfg["agent"] = a.agent
    if getattr(a, "subshell", None) is not None:
        cfg["subshell"] = a.subshell  # --subshell / --no-subshell overrides the config
    # Both of these are answerable from the config alone, so answer them before
    # creating anything — the same rule resolve_agent already follows.
    if a.host and not any(r["name"] == a.host for r in cfg["repos"]):
        known = ", ".join(r["name"] for r in cfg["repos"])
        die(f"--host {a.host!r} is not a configured repo (have: {known})")
    spec = None if a.no_launch else resolve_agent(cfg)  # fail before touching disk

    ours: set[str] = set()  # repos where the branch will be ours to delete on rollback
    missing: list[str] = []  # repos without the branch; only named once ALL are checked
    for r in cfg["repos"]:  # preflight: never create half a set
        # Fetch before looking: --existing usually means a branch someone else
        # pushed, which is a remote ref until we ask. Also lets a base like
        # origin/master resolve below.
        git(r["path"], "fetch", "--quiet", "origin", check=False)  # best effort
        has_branch = git_ok(r["path"], "show-ref", "--verify", "--quiet", f"refs/heads/{branch}")
        # git checks out origin/<branch> into a new local branch on its own, so a
        # remote-only branch counts as existing.
        has_remote = git_ok(r["path"], "show-ref", "--verify", "--quiet", f"refs/remotes/origin/{branch}")
        if not has_branch:
            ours.add(r["name"])  # we are about to create it, so we may delete it
        if a.existing and not (has_branch or has_remote):
            missing.append(r["name"])
        if not a.existing and (has_branch or has_remote):
            if worktrees_of(r["path"]).get(branch):
                die(
                    f"branch {branch!r} already has a worktree in {r['name']} — nothing was created.\n"
                    f"`polytree link {branch}` to launch on it, `polytree rm {branch}` to remove it, "
                    f"or pick another name."
                )
            if has_branch:
                # A lingering local branch with no worktree. `polytree rm` alone keeps
                # an unmerged branch (git branch -d refuses it), so name the command
                # that actually clears it — the plain one is exactly what fails here.
                die(
                    f"branch {branch!r} already exists in {r['name']} (no worktree) — nothing was created.\n"
                    f"`polytree new {branch} --existing` to check it out, `polytree rm {branch} --force` "
                    f"to delete the leftover branch, or pick another name."
                )
            die(
                f"branch {branch!r} already exists on origin in {r['name']} — nothing was created.\n"
                f"`polytree new {branch} --existing` to check it out, or pick another name."
            )
        if worktrees_of(r["path"]).get(branch):
            die(f"{r['name']} already has a worktree for {branch!r} — nothing was created")
        if cfg["backend"] == "git" and Path(dest_for(cfg, r, branch)).exists():
            die(f"{dest_for(cfg, r, branch)} already exists — nothing was created")
        if a.existing:
            continue
        base = configured_base(cfg, r, a.base) or (default_base(r) if cfg["backend"] == "git" else None)
        if base and not git_ok(r["path"], "rev-parse", "--verify", "--quiet", f"{base}^{{commit}}"):
            die(f"base {base!r} does not exist in {r['name']} — nothing was created")

    if missing:
        # Say only what was checked. Which repo trips first in the loop is not
        # evidence about where the branch lives.
        where = ", ".join(missing)
        if a.pr and pr_repo["name"] in missing:
            die(
                f"the branch for PR {a.pr} ({branch!r}) is not on origin in {pr_repo['name']} "
                f"itself — nothing was created.\nIt most likely lives in a fork, which polytree "
                f"cannot fetch yet. Review that PR from Orca, or fetch the branch yourself."
            )
        have = [r["name"] for r in cfg["repos"] if r["name"] not in missing]
        if a.pr:
            die(
                f"PR {a.pr} is on {branch!r}, which is in {', '.join(have)} but not in {where} — "
                f"nothing was created.\nThat pull request does not span every configured repo, "
                f"so there is no set to make. Review it from Orca instead."
            )
        die(
            f"branch {branch!r} exists neither locally nor on origin in {where} — "
            f"nothing was created.\nDrop --existing to create it."
        )

    info(f"Creating '{branch}' worktrees in {len(cfg['repos'])} repos (backend: {cfg['backend']})…")
    created: list[tuple[dict, str]] = []
    try:
        for r in cfg["repos"]:
            path = (create_orca(cfg, r, branch, a.base, a.existing, a.issue) if cfg["backend"] == "orca"
                    else create_git(cfg, r, branch, a.base, a.existing))
            created.append((r, path))
            info(f"  {r['name']:<20} -> {path}")
    except Fail as e:
        rollback(cfg, created, branch, ours)
        tail = f"\n\nRolled back the {len(created)} worktree(s) polytree had created." if created else ""
        die(f"{e}{tail}")
    except BaseException:  # Ctrl-C, OSError… a half-made set is still a half-made set
        rollback(cfg, created, branch, ours)
        raise

    _, host_path = pick_host(created, a.host)
    if a.no_launch:
        record_dir(host_path)  # so the pt() shell wrapper still cds into the set
        return
    launch(cfg, spec, host_path, [p for _, p in created if p != host_path], a.prompt)


def cmd_link(cfg: dict, a: argparse.Namespace) -> None:
    branch = a.branch or current_branch()
    pairs = sibling_map(cfg, branch)
    present = [(r, p) for r, p in pairs if p]
    missing = [r["name"] for r, p in pairs if not p]
    if len(present) < 2:
        # Suggest the command that will actually work: if the branch already
        # exists (a colleague's PR), plain `new` would just refuse.
        exists_somewhere = any(
            branch_exists(r, branch)
            or git_ok(r["path"], "show-ref", "--verify", "--quiet", f"refs/remotes/origin/{branch}")
            for r, _ in pairs
        )
        fix = f"polytree new {branch} --existing" if exists_somewhere else f"polytree new {branch}"
        die(
            f"need >=2 worktrees for branch {branch!r}, found {len(present)}"
            + (f" (missing: {', '.join(missing)})" if missing else "")
            + f".\nCreate the set with: {fix}"
        )
    if missing:
        info(f"warning: no worktree for {branch!r} in: {', '.join(missing)}")
    # Never silently promote another repo to host: hooks/settings come from the
    # host only, so an implicit shift would quietly change what the agent loads.
    if not a.host and present[0][0] is not cfg["repos"][0]:
        die(
            f"the host repo ({cfg['repos'][0]['name']}) has no worktree for {branch!r}.\n"
            f"Pass --host <repo> to run the agent elsewhere — its hooks and settings "
            f"are the ones that will load."
        )
    if a.agent:
        cfg["agent"] = a.agent
    if getattr(a, "subshell", None) is not None:
        cfg["subshell"] = a.subshell  # --subshell / --no-subshell overrides the config
    spec = resolve_agent(cfg)
    _, host_path = pick_host(present, a.host)
    launch(cfg, spec, host_path, [p for _, p in present if p != host_path], a.prompt)


def cmd_paths(cfg: dict, a: argparse.Namespace) -> None:
    for _, p in sibling_map(cfg, a.branch or current_branch()):
        if p:
            info(p)


def cmd_rm(cfg: dict, a: argparse.Namespace) -> None:
    a.branch = a.branch or current_branch()  # no branch given -> the set you stand in
    try:  # captured now: the removals below may delete the dir out from under us
        here: Path | None = Path.cwd()
    except OSError:
        here = None
    # A repo counts if it has a live worktree OR just the branch: someone may have
    # deleted a worktree by hand, or a rollback kept the branch. Skipping those
    # silently left the branch and a stale registration behind, and then `new`
    # refused while pointing at this very command, which found nothing. Wedged.
    present: list[tuple[dict, str | None]] = []
    for r in cfg["repos"]:
        path = worktrees_of(r["path"]).get(a.branch)
        if path or branch_exists(r, a.branch):
            present.append((r, path))
    if not present:
        die(f"nothing to remove for branch {a.branch!r}: no worktree and no branch")

    # Preflight every repo: removing repo 1 and then bailing on repo 2 is the
    # same half-a-set failure `new` refuses to make.
    blocked = [(r["name"], b) for r, p in present if p and (b := removal_blocker(r, p, a.force))]
    if blocked:
        detail = ", ".join(f"{n} ({why})" for n, why in blocked)
        hint = "" if a.force else "\nCommit or stash them, or re-run with --force to discard them."
        die(f"cannot remove: {detail} — nothing was removed.{hint}")

    # Remove the worktree we are standing in LAST. On the orca backend, deleting
    # it tears down this very terminal — Orca kills the process the moment its
    # worktree vanishes — so anything still on the list would be stranded. Doing
    # our own repo last means every sibling is already gone by the time we go.
    if here:
        present.sort(key=lambda rp: bool(rp[1]) and (Path(rp[1]) == here or Path(rp[1]) in here.parents))

    # Confirm before destroying anything. polytree links siblings purely by branch
    # name, so a same-named worktree created independently in another repo looks
    # exactly like part of the set — showing the list is the only way to catch that
    # coincidence before it is removed. --yes skips this for scripts.
    if not getattr(a, "yes", False):
        info(f"About to remove branch {a.branch!r}:")
        for repo, path in present:
            info(f"  {repo['name']:<20} {path or '(branch only, no worktree)'}")
        if a.force:
            info("  --force: discards uncommitted changes and deletes the branch even if unmerged.")
        if not sys.stdin.isatty():
            die("not a terminal — pass --yes to remove without confirming")
        try:
            answer = input("Continue? [y/N] ").strip().lower()
        except EOFError:
            answer = ""
        if answer not in ("y", "yes"):
            info("aborted — nothing was removed")
            return

    done: list[str] = []
    kept: list[dict] = []
    for repo, path in present:
        # `orca worktree rm` deletes the branch too, merged or not. Remember where
        # it pointed so the keep-or-delete decision stays ours and identical on
        # both backends, instead of Orca quietly discarding unmerged work.
        sha = branch_sha(repo, a.branch)
        if path:
            try:
                remove_worktree(cfg, repo, path, force=a.force)
            except Fail as e:
                # removal_blocker approximates git's rules; if git refuses anyway,
                # say exactly what happened instead of leaving the user to guess.
                die(
                    f"{e}\n\nAlready removed: {', '.join(done) or 'nothing'}. "
                    f"Remaining repos untouched.\nRe-run with --force, or clean up by hand."
                )
        else:
            # No live worktree: drop whatever stale registration git still holds.
            git(repo["path"], "worktree", "prune", check=False)
        if sha and not branch_exists(repo, a.branch):
            git(repo["path"], "branch", a.branch, sha, check=False)  # restore, then decide
        git_ok(repo["path"], "branch", "-D" if a.force else "-d", a.branch)
        # Report what is actually true now, not what the delete's exit code implied.
        if branch_exists(repo, a.branch):
            kept.append(repo)
        note = "  (branch kept: not merged)" if repo in kept else ""
        what = f"removed {path}" if path else "no worktree (branch only)"
        info(f"  {repo['name']:<20} {what}{note}")
        done.append(repo["name"])

    if kept:
        # Not "re-run with --force": the worktrees are gone now, so `rm` would only
        # say it found nothing. Give the command that actually works.
        info(f"\nKept unmerged branch {a.branch!r}. To delete it anyway:")
        for repo in kept:
            info(f"  git -C {repo['path']} branch -D {a.branch}")

    # If we just deleted the directory the caller is standing in, their shell is
    # now in a path that no longer exists. Say so — git worktree remove is silent
    # about it, and the confusing symptom is the next command failing in the void.
    if here and any(p and (Path(p) == here or Path(p) in here.parents) for _, p in present):
        info("\n(You were in a removed worktree — cd elsewhere.)")


def cmd_list(cfg: dict, _a: argparse.Namespace) -> None:
    """Show the feature sets: branches with a worktree in 2+ configured repos.

    This is the answer to "what am I working on?" — `link` needs a branch name,
    and nobody remembers those.
    """
    sets: dict[str, list[tuple[dict, str]]] = {}
    for r in cfg["repos"]:
        for branch, path in worktrees_of(r["path"]).items():
            if path != r["path"]:  # a main checkout is not part of a feature set
                sets.setdefault(branch, []).append((r, path))
    complete = {b: v for b, v in sets.items() if len(v) >= 2}
    if not complete:
        info("no feature sets (no branch has a worktree in 2+ configured repos)")
        return
    for branch in sorted(complete):
        entries = complete[branch]
        info(f"{branch}  ({len(entries)}/{len(cfg['repos'])} repos)")
        for repo, path in entries:
            state = "dirty" if is_dirty(path) else "clean"
            info(f"  {repo['name']:<20} {state:<6} {path}")


def cmd_ls(cfg: dict, _a: argparse.Namespace) -> None:
    info(f"backend : {cfg['backend']}" + ("" if cfg["backend"] == "git" else f"  ({orca_bin()})"))
    info(f"agent   : {cfg['agent']}")
    if cfg["backend"] == "git":
        info(f"root    : {cfg['root']}")
    info("repos   : (host first)")
    for i, r in enumerate(cfg["repos"]):
        info(f"  {'* ' if i == 0 else '  '}{r['name']:<20} {r['path']}")


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="polytree",
        description="One feature, many repos: the same branch everywhere, one agent that sees it all.",
        epilog=EXAMPLE,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("--version", action="version", version=f"polytree {VERSION}")
    # Shared by every subcommand, so `-v` works after the command too (polytree new -v …).
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("-v", "--verbose", action="store_true",
                        help="echo each git/orca command as it runs (to stderr)")
    sub = p.add_subparsers(dest="cmd", metavar="<command>")

    n = sub.add_parser("new", parents=[common], help="create the worktree set in every repo + launch the agent")
    n.add_argument("name", nargs="?", help="branch name, used in every repo")
    n.add_argument("--host", metavar="REPO", help="repo that hosts the agent (default: first in config)")
    n.add_argument("--base", metavar="REF", help="branch off this ref instead of the configured base (e.g. --base origin/master for a hotfix)")
    n.add_argument("--pr", metavar="[REPO#]N",
                   help="take the branch from a GitHub PR (implies --existing). PR numbers are "
                        "per-repo: qualify as <repo>#<n>, or plain <n> for the host repo")
    n.add_argument("--existing", action="store_true", help="check out a branch that already exists")
    n.add_argument("--agent", metavar="NAME", help="agent for this run, overriding the config")
    n.add_argument("--issue", metavar="N", help="link the worktrees to a GitHub issue (orca backend)")
    n.add_argument("--prompt", metavar="TEXT", help="initial prompt for the agent")
    n.add_argument("--no-launch", action="store_true", help="create the worktrees, don't start the agent")
    n.add_argument("--subshell", action=argparse.BooleanOptionalAction, default=None,
                   help="git backend: on agent exit, land in a shell in the worktree (default: on)")
    n.set_defaults(fn=cmd_new)

    lk = sub.add_parser("link", parents=[common], help="existing worktrees: launch the agent across them")
    lk.add_argument("branch", nargs="?", help="default: the branch of the current directory")
    lk.add_argument("--host", metavar="REPO", help="repo that hosts the agent (default: first in config)")
    lk.add_argument("--agent", metavar="NAME", help="agent for this run, overriding the config")
    lk.add_argument("--prompt", metavar="TEXT", help="initial prompt for the agent")
    lk.add_argument("--subshell", action=argparse.BooleanOptionalAction, default=None,
                    help="git backend: on agent exit, land in a shell in the worktree (default: on)")
    lk.set_defaults(fn=cmd_link)

    pa = sub.add_parser("paths", parents=[common], help="print the sibling worktree paths (no side effects)")
    pa.add_argument("branch", nargs="?", help="default: the branch of the current directory")
    pa.set_defaults(fn=cmd_paths)

    rm = sub.add_parser("rm", parents=[common], help="remove the worktree set for a branch (refuses if dirty/locked)")
    rm.add_argument("branch", nargs="?", help="default: the branch of the current directory")
    rm.add_argument("--force", action="store_true",
                    help="discard uncommitted changes, remove locked worktrees, and delete the "
                         "branch even if unmerged")
    rm.add_argument("-y", "--yes", action="store_true",
                    help="skip the confirmation prompt (for scripts)")
    rm.set_defaults(fn=cmd_rm)

    li = sub.add_parser("list", parents=[common], help="show your feature sets: branches with a worktree in 2+ repos")
    li.set_defaults(fn=cmd_list)

    ls = sub.add_parser("ls", parents=[common], help="show the resolved config")
    ls.set_defaults(fn=cmd_ls)
    return p


def main() -> None:
    try:
        sys.stdout.reconfigure(line_buffering=True)  # keep our output ordered with git's stderr
    except Exception:
        pass
    parser = build_parser()
    args = parser.parse_args()
    if getattr(args, "verbose", False):
        global VERBOSE
        VERBOSE = True
    if not getattr(args, "fn", None):
        parser.print_help()
        return
    try:
        args.fn(load_config(), args)
    except Fail as e:
        print(f"polytree: {e}", file=sys.stderr)
        sys.exit(1)
    except KeyboardInterrupt:
        sys.exit(130)
    except BrokenPipeError:
        # `polytree list | head` closes the pipe on us. Exit like every other CLI
        # does, rather than dumping a traceback; the devnull dance stops Python
        # from complaining again while flushing at exit.
        os.dup2(os.open(os.devnull, os.O_WRONLY), sys.stdout.fileno())
        sys.exit(141)  # 128 + SIGPIPE


if __name__ == "__main__":
    main()
