tools/devshell/worktree/main.py

#!/usr/bin/env fuchsia-vendored-python
# Copyright 2026 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import sys

from subcommands import add as add_cmd
from subcommands import list as list_cmd
from subcommands import locate as locate_cmd
from subcommands import metrics as metrics_cmd
from subcommands import pool_add as pool_add_cmd
from subcommands import pool_list as pool_list_cmd
from subcommands import pool_remove as pool_remove_cmd
from subcommands import remove as remove_cmd
from worktree import WorktreeState
from worktree_metrics import WorktreeMetricsTracker
from worktree_pool import WorktreePool


def main() -> None:
    parser = argparse.ArgumentParser(
        prog="fx worktree",
        description="Manage Fuchsia worktrees for parallel development.",
        epilog=(
            "Workflow:\n"
            "  • fx worktree add <name>     Create a new worktree.\n"
            "  • fx worktree remove <name>  Remove a worktree.\n"
            "  • fx worktree metrics        Display worktree usage metrics.\n"
            "\n"
            "Advanced (Internal Pool Management):\n"
            "  For performance, worktrees are transparently backed by a pool of reusable checkouts.\n"
            "  You can manage this pool using 'fx worktree pool' commands (e.g., 'fx worktree pool list')."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    subparsers = parser.add_subparsers(
        dest="subcommand", required=True, metavar="COMMAND"
    )

    # Subcommand group 'pool'
    parser_pool = subparsers.add_parser(
        "pool",
        description="Administrative commands for provisioning and maintaining physical checkouts on disk.",
    )
    pool_subparsers = parser_pool.add_subparsers(
        dest="pool_subcommand", required=True
    )

    parser_pool_list = pool_subparsers.add_parser(
        "list",
        help="List all physical worktrees in the pool",
        description="List all physical worktrees in the pool along with state and physical paths.",
    )
    parser_pool_list.add_argument(
        "--json", action="store_true", help="Output JSON"
    )

    parser_pool_add = pool_subparsers.add_parser(
        "add",
        help="Add a new physical worktree checkout to the pool",
    )
    parser_pool_add.add_argument(
        "name", nargs="?", help="Optional physical name"
    )
    parser_pool_add.add_argument("--set", action="append", help="Run 'fx set'")
    group = parser_pool_add.add_mutually_exclusive_group()
    group.add_argument(
        "--symlink-local",
        action="store_true",
        help="Symlink the 'local' directory from the main checkout",
    )
    group.add_argument(
        "--copy-local",
        action="store_true",
        help="Copy the 'local' directory from the main checkout",
    )

    parser_pool_remove = pool_subparsers.add_parser(
        "remove",
        help="Remove a physical worktree from the pool",
    )
    parser_pool_remove.add_argument("name", help="Name of physical worktree")
    parser_pool_remove.add_argument(
        "--force", action="store_true", help="Force removal"
    )

    # Subcommand 'locate'
    parser_locate = subparsers.add_parser(
        "locate",
        help="Print the absolute path to a worktree directory",
        description="Locate a worktree by name and print its absolute path on disk.",
    )
    parser_locate.add_argument("name", help="Name of the worktree")

    # Subcommand 'list'
    parser_list = subparsers.add_parser(
        "list",
        help="List all leased worktrees",
        description="List all leased worktrees and their git branches.",
    )
    parser_list.add_argument("--json", action="store_true", help="Output JSON")

    # Subcommand 'metrics'
    parser_metrics = subparsers.add_parser(
        "metrics",
        help="Display local worktree usage metrics and statistics",
        description="Display time-weighted average, peak, and minimum worktree usage statistics.",
    )
    parser_metrics.add_argument(
        "--json", action="store_true", help="Output raw JSON metrics"
    )

    # Subcommand 'add'
    parser_add = subparsers.add_parser(
        "add",
        help="Add a leased worktree checkout for development",
    )
    parser_add.add_argument("name", help="Name of worktree / branch")
    parser_add.add_argument(
        "--pool-name", help="Specific pool slot to allocate"
    )
    parser_add.add_argument("--json", action="store_true", help="Output JSON")

    # Subcommand 'remove'
    parser_remove = subparsers.add_parser(
        "remove",
        help="Remove a leased worktree and return it to the pool",
    )
    parser_remove.add_argument("name", help="Name of worktree to remove")

    # Internal subcommand '_complete' for shell completion
    parser_complete = subparsers.add_parser(
        "_complete",
        help="Internal completion helper",
    )
    parser_complete.add_argument(
        "filter",
        choices=[
            "physical_free",
            "physical_leased",
            "physical_all",
            "leased",
            "all",
        ],
    )

    args = parser.parse_args()
    pool = WorktreePool()

    try:
        if args.subcommand == "pool":
            if args.pool_subcommand == "list":
                pool_list_cmd.run(args, pool)
            elif args.pool_subcommand == "add":
                pool_add_cmd.run(args, pool)
            elif args.pool_subcommand == "remove":
                pool_remove_cmd.run(args, pool)
        elif args.subcommand == "locate":
            locate_cmd.run(args, pool)
        elif args.subcommand == "list":
            list_cmd.run(args, pool)
        elif args.subcommand == "metrics":
            metrics_cmd.run(args, pool)
        elif args.subcommand == "add":
            add_cmd.run(args, pool)
        elif args.subcommand == "remove":
            remove_cmd.run(args, pool)
        elif args.subcommand == "_complete":
            print(" ".join(pool.get_completion_choices(args.filter)))
        else:
            print(f"Unknown subcommand: {args.subcommand}", file=sys.stderr)
            sys.exit(1)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    finally:
        _record_metrics_safely(pool)


def _record_metrics_safely(pool: WorktreePool) -> None:
    try:
        worktrees = pool.get_worktrees()
        total_count = len(worktrees)
        leased_count = sum(
            1 for wt in worktrees if wt.get_state() == WorktreeState.LEASED
        )
        built_recently_count = sum(
            1 for wt in worktrees if wt.is_built_recently(max_age_days=7)
        )
        not_built_recently_count = total_count - built_recently_count
        tracker = WorktreeMetricsTracker(pool.jiri_root, pool.fuchsia_dir)
        tracker.record_state(
            current_total=total_count,
            current_leased=leased_count,
            current_built_recently=built_recently_count,
            current_not_built_recently=not_built_recently_count,
        )
    except Exception as e:
        print(
            f"Warning: Failed to record worktree metrics: {e}",
            file=sys.stderr,
        )


if __name__ == "__main__":
    main()