#!python
#
# Optional dependencies:
#   - colorama
#   - setproctitle
#
# Author: James Cherti
# URL: https://github.com/jamescherti/git-rexec
# Version: 1.0.4
#
# Copyright (C) 2019-2026 James Cherti
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.
#
"""Find Git repositories and execute commands against them in parallel.

This script allows for finding Git repositories within a directory structure
and executing commands against them. It supports conditional filtering,
background execution, and sequential foreground execution.
"""

from __future__ import annotations

import argparse
import os
import shlex
import shutil
import subprocess
import sys
import textwrap
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from multiprocessing import cpu_count
from pathlib import Path

try:
    from colorama import Fore
    from colorama import init as colorama_init

    HAS_COLORAMA: bool = True
except ImportError:
    HAS_COLORAMA = False

    class Fore:  # type: ignore
        """Provide fallback empty strings for missing colorama Fore attributes."""

        YELLOW: str = ""
        RED: str = ""
        RESET: str = ""


@dataclass
class CommandResult:
    """Store the result of a subprocess command execution.

    :param command: The command that was executed.
    :param returncode: The exit status of the command.
    :param stdout: Standard output content.
    :param stderr: Standard error content.
    """

    command: list[str]
    returncode: int
    stdout: str = ""
    stderr: str = ""


@dataclass
class RepoContext:
    """Represent a processed Git repository and its execution states.

    :param path: The absolute path to the repository root.
    :param parallel_result: The result of the parallel execution, if any.
    """

    path: Path
    parallel_result: CommandResult | None = None


def git_toplevel(repo_path: Path) -> Path | None:
    """Return the absolute path to the top-level directory of a Git repository.

    :param repo_path: Path inside the Git repository.
    :return: Path to the repository's top-level directory, or None if not a
             Git repo.
    """
    try:
        proc: subprocess.CompletedProcess[str] = \
            subprocess.run(["git", "-C", str(repo_path), "rev-parse",
                            "--show-toplevel"],
                           capture_output=True,
                           text=True,
                           check=True)
        return Path(proc.stdout.strip()).resolve()
    except subprocess.CalledProcessError:
        return None


def git_find_repo(root: Path, exclude_paths: list[Path] | None = None) \
        -> tuple[set[Path], set[Path]]:
    """Recursively find all Git repositories under the root path using Python.

    :param root: Directory to search under.
    :param exclude_paths: Directories to exclude to optimize traversal.
    :return: A tuple containing standard repository paths and sub-repository
    paths.
    """
    root_res: Path = root.resolve()

    # Paths are already resolved by discover_repos
    ex_resolved: set[Path] = set(exclude_paths) if exclude_paths else set()

    # Do not begin traversal if the root itself is excluded
    repos: set[Path] = set()
    sub_repos: set[Path] = set()
    if ex_resolved and any((root_res == ex_path or
                            root_res.is_relative_to(ex_path))
                           for ex_path in ex_resolved):
        return repos, sub_repos

    # Store Path objects internally to simplify maintenance
    dirs_to_visit: list[Path] = [root_res]
    while dirs_to_visit:
        current_dir: Path = dirs_to_visit.pop()

        try:
            with os.scandir(current_dir) as iter_scandir:
                dirs_in_current: list[Path] = []
                is_repo: bool = False

                for entry in iter_scandir:
                    if entry.name == ".git":
                        is_repo = True
                        if entry.is_dir():
                            repos.add(current_dir)
                        elif entry.is_file():
                            sub_repos.add(current_dir)

                        # Found the repository, stop evaluating this directory
                        break

                    try:
                        # Using follow_symlinks=False avoids an extra
                        # is_symlink() call
                        if entry.is_dir(follow_symlinks=False):
                            entry_path: Path = Path(entry.path)
                            if ex_resolved and entry_path in ex_resolved:
                                continue
                            dirs_in_current.append(entry_path)
                    except OSError:
                        continue

                if not is_repo:
                    # Append collected directories only if we are not inside a
                    # repository
                    dirs_to_visit.extend(dirs_in_current)
        except OSError:
            continue

    return repos, sub_repos


def get_sub_repos(repo_path: Path) -> set[Path]:
    """Extract worktrees and submodules from a Git repository.

    :param repo_path: The path to the standard repository.
    :return: A set of absolute paths to worktrees and submodules.
    """
    sub_repos: set[Path] = set()
    repo_path_resolved: Path = repo_path.resolve()

    # Discover worktrees
    # pylint: disable=too-many-try-statements
    try:
        wt_proc: subprocess.CompletedProcess[str] = \
            subprocess.run(["git", "-C", str(repo_path),
                            "worktree", "list", "--porcelain"],
                           capture_output=True,
                           text=True,
                           check=True)

        current_wt: Path | None = None
        is_bare: bool = False

        for line in wt_proc.stdout.splitlines():
            line = line.strip()

            # Git porcelain separates blocks with a blank line
            if not line:
                if (current_wt and
                        not is_bare and current_wt != repo_path_resolved):
                    sub_repos.add(current_wt)
                # Reset state for the next block
                current_wt = None
                is_bare = False
                continue

            if line.startswith("worktree "):
                current_wt = Path(line[9:]).resolve()
            elif line == "bare":
                is_bare = True

        # Process the final block in case the output lacks a trailing newline
        if current_wt and not is_bare and current_wt != repo_path_resolved:
            sub_repos.add(current_wt)

    except subprocess.CalledProcessError:
        pass

    # Discover submodules
    try:
        sm_proc: subprocess.CompletedProcess[str] = \
            subprocess.run(["git", "-C", str(repo_path), "submodule",
                            "foreach", "--quiet", "pwd",],
                           capture_output=True,
                           text=True,
                           check=True)

        for line in sm_proc.stdout.splitlines():
            sm_path: Path = Path(line.strip()).resolve()
            sub_repos.add(sm_path)

    except subprocess.CalledProcessError:
        pass

    return sub_repos


def run_command_get_output(repo_path: Path,
                           cmd_list: list[str],
                           capture: bool = True) -> CommandResult:
    """Execute a shell command within a specific directory.

    :param repo_path: The directory in which to execute the command.
    :param cmd_list: The command list to execute.
    :param capture: Whether to capture stdout/stderr.
    :return: The result of the command execution.
    """
    if not cmd_list:
        return CommandResult(command=[], returncode=0)

    try:
        proc: subprocess.CompletedProcess[str] = \
            subprocess.run(cmd_list, cwd=repo_path, capture_output=capture,
                           text=True, check=False,)
        return CommandResult(command=cmd_list,
                             returncode=proc.returncode,
                             stdout=proc.stdout if proc.stdout else "",
                             stderr=proc.stderr if proc.stderr else "")
    except FileNotFoundError:
        return CommandResult(
            command=cmd_list,
            returncode=127,
            stderr=f"Error: Command not found: '{cmd_list[0]}'\n"
        )


def format_parallel_output(repo_path: Path,
                           result: CommandResult, quiet: bool = False) -> str:
    """Format the output of a parallel command for display.

    :param repo_path: The repository path.
    :param result: The execution result.
    :param quiet: Suppress informational headers.
    :return: Formatted string ready for printing.
    """
    raw_output: str = result.stdout + result.stderr
    if not raw_output and result.returncode == 0:
        return ""

    formatted: str = raw_output.replace("\t", "    ").rstrip()
    if formatted:
        formatted += "\n"

    if quiet:
        return formatted

    header: str = f"{Fore.YELLOW}[EXEC-P] {repo_path}"
    cmd_str: str = " ".join(result.command)
    header += f": {cmd_str}{Fore.RESET}\n"

    indent: str = " " * 4
    indented_body: str = textwrap.indent(formatted, prefix=indent)
    return header + indented_body


def process_repo(repo_path: Path,
                 exec_parallel_cmd: list[str] | None,
                 if_exec_cmd: list[str] | None) -> RepoContext | None:
    """Process a single repository: check conditions and run parallel commands.

    :param repo_path: The repository path.
    :param exec_parallel_cmd: Command list to execute in background/parallel.
    :param if_exec_cmd: Command list to check before processing (filter).
    :return: RepoContext if processed successfully, None if filtered out.
    """
    # Filter: if-exec
    if_exec_cmd_list: list[str] | None = if_exec_cmd
    if if_exec_cmd_list:
        # We discard output for the filter check, only caring about exit code
        filter_res: CommandResult = run_command_get_output(repo_path,
                                                           if_exec_cmd_list,
                                                           capture=True)
        if filter_res.returncode != 0:
            return None

    # Action: exec-parallel
    parallel_result: CommandResult | None = None
    exec_parallel_cmd_list: list[str] | None = exec_parallel_cmd
    if exec_parallel_cmd_list:
        parallel_result = \
            run_command_get_output(repo_path,
                                   exec_parallel_cmd_list,
                                   capture=True)

    return RepoContext(path=repo_path, parallel_result=parallel_result)


def discover_repos(directory: Path,
                   max_workers: int,
                   exclude_dirs: list[str],
                   include_sub_repos: bool) -> set[Path]:
    """Discover Git repositories starting from directory and apply exclusions.

    :param directory: The root directory for search.
    :param max_workers: Maximum number of threads/workers.
    :param exclude_dirs: Directories to exclude.
    :param include_sub_repos: Whether to include sub-repositories.
    :return: A set of repository root paths.
    """
    repos: set[Path] = set()
    worktrees: set[Path] = set()

    list_ex_paths: list[Path] = (
        [Path(ex_dir).resolve() for ex_dir in exclude_dirs]
        if exclude_dirs
        else []
    )

    # Discover all repositories
    toplevel: Path | None = git_toplevel(directory)
    if toplevel:
        repos = {toplevel}
    else:
        repos, worktrees = git_find_repo(directory, list_ex_paths)

    # Extract and append sub-repositories concurrently if requested
    if include_sub_repos:
        sub_repos: set[Path] = set()
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures: list[Future[set[Path]]] = [
                executor.submit(get_sub_repos, repo)
                for repo in repos | worktrees
            ]
            for future in as_completed(futures):
                result: set[Path] = future.result()
                sub_repos.update(result)

        repos.update(sub_repos)

    # Filter out excluded directories (applies dynamically to sub-repos too)
    if not list_ex_paths:
        return repos

    return {
        repo
        for repo in repos
        if not any(repo.is_relative_to(ex_path)
                   for ex_path in list_ex_paths)
    }


def execute_parallel_tasks(repos: set[Path],
                           exec_parallel: list[str] | None,
                           if_exec: list[str] | None,
                           max_workers: int) -> list[RepoContext]:
    """Run discovery and parallel execution tasks.

    :param repos: set of repositories to process.
    :param exec_parallel: Command list for background execution.
    :param if_exec: Command list for conditional filtering.
    :param max_workers: Maximum number of threads for execution.
    :return: list of processed repository contexts.
    """
    results: list[RepoContext] = []

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures: dict[Future[RepoContext | None], Path] = {
            executor.submit(process_repo, repo, exec_parallel, if_exec): repo
            for repo in repos
        }

        for future in as_completed(futures):
            repo_path: Path = futures[future]
            try:
                result: RepoContext | None = future.result()
                if result:
                    results.append(result)
            except (OSError, ValueError) as exc:
                print(f"Error processing repository {repo_path}: {exc}",
                      file=sys.stderr)

    return results


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments.

    :return: Parsed arguments.
    """
    parser: argparse.ArgumentParser = argparse.ArgumentParser(
        description=("Find Git repositories and execute commands against "
                     "them in parallel."),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "-C",
        "--directory",
        type=Path,
        default=Path("."),
        help="Root directory to search (defaults to current directory)",
    )
    parser.add_argument(
        "--exclude-dir",
        action="append",
        default=[],
        help="Exclude a specific directory and all of its subdirectories",
    )
    parser.add_argument(
        "-p",
        "--parallel",
        action="store_true",
        help="Execute the command in parallel using threads",
        default=False,
    )
    parser.add_argument(
        "-i",
        "--if-exec",
        type=str,
        help="Execute commands only if this check returns exit code 0.",
        default=None,
    )
    parser.add_argument(
        "-j",
        "--jobs",
        type=int,
        dest="max_workers",
        help="Maximum number of processors/workers to use",
        default=(cpu_count() or 1),
    )
    parser.add_argument(
        "-q",
        "--quiet",
        action="store_true",
        help=(
            "Quiet mode. Suppresses the informational log prefixes "
            "([EXEC] and [EXEC-P]) that precede execution output."
        ),
        default=False,
    )
    parser.add_argument(
        "-s",
        "--include-sub-repos",
        action="store_true",
        help="Include sub-repositories (e.g., Git worktrees and submodules)",
        default=False,
    )
    parser.add_argument(
        "--print",
        action="store_true",
        help="Print the paths (only when no command is provided)",
        default=False,
    )
    parser.add_argument(
        "--print0",
        action="store_true",
        help=(
            "Separate the paths with a null character (only when no "
            "command is provided)"
        ),
        default=False,
    )
    parser.add_argument(
        "exec_cmd",
        type=str,
        nargs="*",
        help="The command to execute. You can use -- to pass options.",
    )

    args: argparse.Namespace = parser.parse_args()

    return args


def print_error_summary(errors: list[tuple[Path, CommandResult]]) -> int:
    """Print a summary of execution errors.

    :param errors: list of tuples containing path and result.
    :return: The final exit code (1 if errors exist, else 0).
    """
    if not errors:
        return 0

    print()
    print(f"{Fore.RED}Errors:{Fore.RESET}")

    final_errno: int = 0
    for repo_path, result in errors:
        cmd_display: str = " ".join(result.command)
        if result.returncode != 0:
            final_errno = 1
            if result.returncode == 127:
                msg: str = "Command not found"
            else:
                msg = f"errno {result.returncode}"

            print(
                f"{Fore.RED}  - {repo_path}: {msg}: "
                f"{cmd_display}{Fore.RESET}"
            )

    return final_errno


def main() -> None:
    """Execute the main command-line interface."""
    # Optional: setproctitle
    try:
        # pylint: disable=import-outside-toplevel
        from setproctitle import setproctitle

        setproctitle(Path(sys.argv[0]).name)  # type: ignore
    except ImportError:
        # Optional dependency 'setproctitle' is not installed.
        pass

    if HAS_COLORAMA:
        colorama_init()

    # Disable git prompting
    os.environ["GIT_TERMINAL_PROMPT"] = "0"

    args: argparse.Namespace = parse_args()

    # Check for empty action
    if not args.exec_cmd and not args.print and not args.print0:
        print("Nothing to do.", file=sys.stderr)
        sys.exit(2)

    # Verify that the execution command exists
    if args.exec_cmd and not shutil.which(args.exec_cmd[0]):
        print(f"Error: Command not found: '{args.exec_cmd[0]}'",
              file=sys.stderr)
        sys.exit(127)

    # Discover Repositories
    repos: set[Path] = discover_repos(args.directory.absolute(),
                                      args.max_workers,
                                      args.exclude_dir,
                                      args.include_sub_repos)

    # setup background command list if needed
    exec_parallel_cmd: list[str] | None = None
    if args.parallel and args.exec_cmd:
        exec_parallel_cmd = args.exec_cmd

    if_exec_cmd: list[str] | None = None
    if args.if_exec:
        if_exec_cmd = shlex.split(args.if_exec)

    # Parallel Processing (Filter + Background Exec)
    processed_repos: list[RepoContext] = \
        execute_parallel_tasks(repos,
                               exec_parallel_cmd,
                               if_exec_cmd,
                               args.max_workers)

    execution_errors: list[tuple[Path, CommandResult]] = []

    # Main Loop: Display results and run sequential commands
    for context in processed_repos:
        repo_path: Path = context.path

        # Handle Background Execution Results
        if context.parallel_result:
            output_display: str = \
                format_parallel_output(repo_path,
                                       context.parallel_result,
                                       args.quiet)
            if output_display:
                print(output_display, end="")

            if context.parallel_result.returncode != 0:
                execution_errors.append((repo_path, context.parallel_result))

        # Handle Foreground Execution
        if args.exec_cmd and not args.parallel:
            if not args.quiet:
                print(f"{Fore.YELLOW}[EXEC] {repo_path}: "
                      f"{shlex.join(args.exec_cmd)}"
                      f"{Fore.RESET}")

            # Run interactively/sequentially (capture=False allows interaction)
            # However, for consistency with error tracking, we might want to
            # capture. Usually --exec implies seeing output immediately.
            try:
                # We use subprocess.check_call to allow direct stdout/stderr
                # flow unless we want to capture for error summary. To match
                # previous logic, we let it flow to stdout.
                subprocess.check_call(args.exec_cmd, cwd=repo_path)
            except subprocess.CalledProcessError as err:
                # Construct a dummy result for the error summary
                execution_errors.append(
                    (repo_path,
                     CommandResult(command=args.exec_cmd,
                                   returncode=err.returncode)))
            except FileNotFoundError:
                execution_errors.append(
                    (repo_path,
                     CommandResult(command=args.exec_cmd,
                                   returncode=127)))

        # If no commands were run, just list the repo
        if not args.exec_cmd:
            if args.print0:
                print(repo_path, end="\0")
            elif args.print:
                print(repo_path)

    # Final Error Summary
    sys.exit(print_error_summary(execution_errors))


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)
    except BrokenPipeError:
        sys.exit(141)
