cli.py

#!/usr/bin/env python3
"""
git-pr-stack: Create and manage stacked PRs on GitHub with zero config.

Usage:
    git pr-stack init
    git pr-stack add <branch>
    git pr-stack remove <branch>
    git pr-stack list
    git pr-stack create
    git pr-stack rebase <onto>
    git pr-stack merge
    git pr-stack sync
    git pr-stack status
"""

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Optional


STACK_FILE = ".git-pr-stack.json"


def get_git_root() -> Path:
    """Get the root directory of the current git repository."""
    result = subprocess.run(
        ["git", "rev-parse", "--show-toplevel"],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        print("❌ Not a git repository", file=sys.stderr)
        sys.exit(1)
    return Path(result.stdout.strip())


def get_current_branch() -> str:
    """Get the name of the current git branch."""
    result = subprocess.run(
        ["git", "rev-parse", "--abbrev-ref", "HEAD"],
        capture_output=True,
        text=True,
    )
    return result.stdout.strip()


def load_stack(git_root: Path) -> dict:
    """Load the stack configuration from the repo."""
    stack_file = git_root / STACK_FILE
    if not stack_file.exists():
        return {"stack": [], "prs": {}}
    with open(stack_file) as f:
        return json.load(f)


def save_stack(git_root: Path, stack: dict) -> None:
    """Save the stack configuration to the repo."""
    stack_file = git_root / STACK_FILE
    with open(stack_file, "w") as f:
        json.dump(stack, f, indent=2)


def run_cmd(cmd: list[str], check: bool = True) -> subprocess.CompletedProcess:
    """Run a shell command and return the result."""
    return subprocess.run(cmd, capture_output=True, text=True, check=check)


def cmd_init(args: argparse.Namespace) -> None:
    """Initialize a new PR stack in the current repository."""
    git_root = get_git_root()
    stack_file = git_root / STACK_FILE

    if stack_file.exists():
        print("⚠️  PR stack already initialized. Use `git pr-stack list` to see it.")
        return

    current = get_current_branch()
    stack = {"stack": [current], "prs": {}, "base": "main"}

    save_stack(git_root, stack)

    # Add stack file to gitignore if not already there
    gitignore = git_root / ".gitignore"
    gitignore_entry = STACK_FILE
    if gitignore.exists():
        with open(gitignore) as f:
            if gitignore_entry not in f.read():
                with open(gitignore, "a") as f_append:
                    f_append.write(f"\n# git-pr-stack\n{gitignore_entry}\n")
    else:
        with open(gitignore, "w") as f:
            f.write(f"# git-pr-stack\n{gitignore_entry}\n")

    print(f"✅ PR stack initialized with branch: {current}")
    print(f"   Base branch set to: main")
    print(f"   Run `git pr-stack add <branch>` to add more branches")


def cmd_add(args: argparse.Namespace) -> None:
    """Add a branch to the top of the stack."""
    git_root = get_git_root()
    stack = load_stack(git_root)
    branch = args.branch

    if not stack["stack"]:
        print("❌ No stack initialized. Run `git pr-stack init` first.")
        return

    if branch in stack["stack"]:
        print(f"⚠️  Branch '{branch}' is already in the stack")
        return

    # Check if branch exists
    result = run_cmd(["git", "rev-parse", "--verify", branch], check=False)
    if result.returncode != 0:
        print(f"❌ Branch '{branch}' does not exist")
        return

    stack["stack"].append(branch)
    save_stack(git_root, stack)

    parent = stack["stack"][-2] if len(stack["stack"]) > 1 else stack.get("base", "main")
    print(f"✅ Added '{branch}' to the stack")
    print(f"   Parent: {parent}")
    print(f"   Stack: {' → '.join(stack['stack'])}")


def cmd_remove(args: argparse.Namespace) -> None:
    """Remove a branch from the stack."""
    git_root = get_git_root()
    stack = load_stack(git_root)
    branch = args.branch

    if branch not in stack["stack"]:
        print(f"❌ Branch '{branch}' is not in the stack")
        return

    if branch == stack["stack"][0]:
        print("❌ Cannot remove the base branch of the stack")
        return

    stack["stack"].remove(branch)
    if branch in stack.get("prs", {}):
        del stack["prs"][branch]
    save_stack(git_root, stack)

    print(f"✅ Removed '{branch}' from the stack")
    print(f"   Stack: {' → '.join(stack['stack'])}")


def cmd_list(args: argparse.Namespace) -> None:
    """Display the current stack with PR status."""
    git_root = get_git_root()
    stack = load_stack(git_root)

    if not stack["stack"]:
        print("❌ No stack initialized. Run `git pr-stack init` first.")
        return

    base = stack.get("base", "main")
    print(f"\n📚 PR Stack (base: {base})\n")
    print(f"  {'Branch':<30} {'PR':<10} {'Status':<15}")
    print(f"  {'─' * 30} {'─' * 10} {'─' * 15}")

    for i, branch in enumerate(stack["stack"]):
        pr_info = stack.get("prs", {}).get(branch, {})
        pr_number = pr_info.get("number", "—")
        pr_status = pr_info.get("status", "no PR")
        pr_url = pr_info.get("url", "")

        current = " ← current" if branch == get_current_branch() else ""
        indent = "  " * i
        print(f"  {indent}{branch:<30} #{pr_number:<9} {pr_status}{current}")

        if pr_url:
            print(f"  {indent}{'':30} {pr_url}")

    print(f"\n  Total: {len(stack['stack'])} branches in stack")
    print()


def cmd_create(args: argparse.Namespace) -> None:
    """Create GitHub PRs for all branches in the stack."""
    git_root = get_git_root()
    stack = load_stack(git_root)

    if not stack["stack"]:
        print("❌ No stack initialized. Run `git pr-stack init` first.")
        return

    base = stack.get("base", "main")
    prs = stack.get("prs", {})

    print("🚀 Creating PRs for the stack...\n")

    for i, branch in enumerate(stack["stack"]):
        if branch in prs and prs[branch].get("number"):
            print(f"  ⏭️  '{branch}' already has PR #{prs[branch]['number']}")
            continue

        # The base for this PR is the previous branch (or 'main' for the first)
        if i == 0:
            pr_base = base
        else:
            pr_base = stack["stack"][i - 1]

        # Create PR using gh CLI
        result = run_cmd(
            ["gh", "pr", "create",
             "--head", branch,
             "--base", pr_base,
             "--title", f"[stack] {branch}",
             "--body", f"Part of PR stack. Base: {pr_base}"],
            check=False,
        )

        if result.returncode == 0:
            # Parse PR URL from output
            output = result.stdout.strip()
            pr_url = output.split()[-1] if output else ""
            # Extract PR number from URL
            pr_number = pr_url.split("/")[-1] if pr_url else "?"

            prs[branch] = {
                "number": pr_number,
                "url": pr_url,
                "status": "open",
                "base": pr_base,
            }
            print(f"  ✅ Created PR #{pr_number}: {branch} → {pr_base}")
        else:
            print(f"  ❌ Failed to create PR for '{branch}': {result.stderr.strip()}")

    stack["prs"] = prs
    save_stack(git_root, stack)
    print("\n✨ Done! Use `git pr-stack list` to see the stack.")


def cmd_rebase(args: argparse.Namespace) -> None:
    """Rebase the entire stack onto a branch."""
    git_root = get_git_root()
    stack = load_stack(git_root)
    onto = args.onto

    if not stack["stack"]:
        print("❌ No stack initialized.")
        return

    print(f"🔄 Rebasing stack onto '{onto}'...\n")

    # Rebase each branch in order
    for i, branch in enumerate(stack["stack"]):
        parent = onto if i == 0 else stack["stack"][i - 1]

        # Checkout the branch
        run_cmd(["git", "checkout", branch])

        # Rebase onto the parent
        print(f"  Rebasing '{branch}' onto '{parent}'...")
        result = run_cmd(["git", "rebase", parent], check=False)

        if result.returncode != 0:
            print(f"  ❌ Conflict while rebasing '{branch}'")
            print(f"     Resolve conflicts and run: git rebase --continue")
            print(f"     Then run `git pr-stack rebase {onto}` again")
            sys.exit(1)

    current = get_current_branch()
    print(f"\n✅ Stack rebased onto '{onto}'")
    print(f"   Current branch: {current}")


def cmd_merge(args: argparse.Namespace) -> None:
    """Merge the bottom PR and advance the stack."""
    git_root = get_git_root()
    stack = load_stack(git_root)

    if not stack["stack"]:
        print("❌ No stack initialized.")
        return

    bottom_branch = stack["stack"][0]
    pr_info = stack.get("prs", {}).get(bottom_branch, {})
    pr_number = pr_info.get("number")

    if not pr_number:
        print(f"❌ Bottom branch '{bottom_branch}' has no PR. Run `git pr-stack create` first.")
        return

    # Merge the PR
    print(f"🔀 Merging PR #{pr_number} ({bottom_branch})...")
    result = run_cmd(
        ["gh", "pr", "merge", pr_number, "--squash", "--delete-branch"],
        check=False,
    )

    if result.returncode != 0:
        print(f"❌ Failed to merge PR #{pr_number}: {result.stderr.strip()}")
        return

    # Remove the merged branch from the stack
    stack["stack"].pop(0)
    if bottom_branch in stack.get("prs", {}):
        del stack["prs"][bottom_branch]

    # Retarget the new bottom PR to the base
    if stack["stack"]:
        new_bottom = stack["stack"][0]
        base = stack.get("base", "main")
        new_pr_info = stack.get("prs", {}).get(new_bottom, {})
        new_pr_number = new_pr_info.get("number")

        if new_pr_number:
            print(f"🔄 Retargeting PR #{new_pr_number} ({new_bottom}) → {base}")
            run_cmd(
                ["gh", "pr", "edit", new_pr_number, "--base", base],
                check=False,
            )
            stack["prs"][new_bottom]["base"] = base

    save_stack(git_root, stack)
    print(f"\n✅ Merged and advanced the stack!")


def cmd_sync(args: argparse.Namespace) -> None:
    """Sync stack state with GitHub (detect merged PRs)."""
    git_root = get_git_root()
    stack = load_stack(git_root)
    prs = stack.get("prs", {})

    if not prs:
        print("No PRs to sync.")
        return

    print("🔄 Syncing stack with GitHub...\n")

    merged_branches = []
    for branch, pr_info in list(prs.items()):
        pr_number = pr_info.get("number")
        if not pr_number:
            continue

        result = run_cmd(
            ["gh", "pr", "view", pr_number, "--json", "state,mergedAt"],
            check=False,
        )

        if result.returncode == 0:
            pr_data = json.loads(result.stdout)
            if pr_data.get("state") == "MERGED":
                print(f"  ✅ PR #{pr_number} ({branch}) was merged")
                merged_branches.append(branch)
            else:
                prs[branch]["status"] = pr_data.get("state", "unknown").lower()

    # Remove merged branches from stack
    for branch in merged_branches:
        if branch in stack["stack"]:
            stack["stack"].remove(branch)
        if branch in prs:
            del prs[branch]

    save_stack(git_root, stack)
    print(f"\n✅ Synced! {len(merged_branches)} merged PRs removed from stack.")


def cmd_status(args: argparse.Namespace) -> None:
    """Check CI status for all PRs in the stack."""
    git_root = get_git_root()
    stack = load_stack(git_root)
    prs = stack.get("prs", {})

    if not prs:
        print("No PRs to check.")
        return

    print("🔍 Checking CI status for all PRs...\n")

    for branch, pr_info in prs.items():
        pr_number = pr_info.get("number")
        if not pr_number:
            continue

        result = run_cmd(
            ["gh", "pr", "checks", pr_number],
            check=False,
        )

        if result.returncode == 0:
            print(f"  PR #{pr_number} ({branch}):")
            for line in result.stdout.strip().split("\n"):
                if line.strip():
                    print(f"    {line}")
        else:
            print(f"  PR #{pr_number} ({branch}): Could not fetch checks")

    print()


def main():
    parser = argparse.ArgumentParser(
        prog="git-pr-stack",
        description="Create and manage stacked PRs on GitHub",
    )
    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # init
    subparsers.add_parser("init", help="Initialize a new PR stack")

    # add
    add_parser = subparsers.add_parser("add", help="Add a branch to the stack")
    add_parser.add_argument("branch", help="Branch name to add")

    # remove
    remove_parser = subparsers.add_parser("remove", help="Remove a branch from the stack")
    remove_parser.add_argument("branch", help="Branch name to remove")

    # list
    subparsers.add_parser("list", help="Display the current stack")

    # create
    subparsers.add_parser("create", help="Create GitHub PRs for the stack")

    # rebase
    rebase_parser = subparsers.add_parser("rebase", help="Rebase the entire stack")
    rebase_parser.add_argument("onto", help="Branch to rebase onto")

    # merge
    subparsers.add_parser("merge", help="Merge the bottom PR and advance")

    # sync
    subparsers.add_parser("sync", help="Sync stack state with GitHub")

    # status
    subparsers.add_parser("status", help="Check CI status for all PRs")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return

    commands = {
        "init": cmd_init,
        "add": cmd_add,
        "remove": cmd_remove,
        "list": cmd_list,
        "create": cmd_create,
        "rebase": cmd_rebase,
        "merge": cmd_merge,
        "sync": cmd_sync,
        "status": cmd_status,
    }

    cmd_func = commands.get(args.command)
    if cmd_func:
        cmd_func(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()