#!python
#
# Optional dependencies:
#   - colorama
#   - setproctitle
#
# Author: James Cherti
# URL: https://github.com/jamescherti/git-rexec
# Version: 1.0.2
#
# 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.
"""

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
from typing import Optional

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
        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: Optional[CommandResult] = None


def git_toplevel(repo_path: Path) -> Optional[Path]:
    """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_py_find_repo(root: Path) -> tuple[set[Path], set[Path]]:
    """Recursively find all Git repositories under the root path using Python.

    :param root: Directory to search under.
    :return: A tuple containing standard repository paths and sub-repository
    paths.
    """
    repos: set[Path] = set()
    sub_repos: set[Path] = set()

    for dirpath, dirnames, filenames in os.walk(root):
        if ".git" in dirnames:
            repos.add(Path(dirpath).absolute())
        elif ".git" in filenames:
            sub_repos.add(Path(dirpath).absolute())
        else:
            continue

        # Clear dirnames to prevent os.walk from descending into
        # subdirectories
        dirnames.clear()

    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()

    # 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: Optional[Path] = 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.resolve():
                    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.resolve():
            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: Optional[list[str]],
                 if_exec_cmd: Optional[list[str]]) -> Optional[RepoContext]:
    """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: Optional[list[str]] = 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: Optional[CommandResult] = None
    exec_parallel_cmd_list: Optional[list[str]] = 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()

    # Discover all repositories
    toplevel: Optional[Path] = git_toplevel(directory)
    if toplevel:
        repos = {toplevel}
    else:
        repos, worktrees = git_py_find_repo(directory)

    # 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
    if not exclude_dirs:
        return repos

    list_ex_paths: list[Path] = [Path(ex_dir).resolve()
                                 for ex_dir in exclude_dirs]
    return {
        repo for repo in repos
        if not any(repo.resolve().is_relative_to(ex_path)
                   for ex_path in list_ex_paths)
    }


def execute_parallel_tasks(
    repos: set[Path],
    exec_parallel: Optional[list[str]],
    if_exec: Optional[list[str]],
    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[Optional[RepoContext]], 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: Optional[RepoContext] = 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 repos and execute commands.",
        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 tracking headers "
              "([EXEC] and [EXEC-P]) that prefix 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(0)

    # 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: Optional[list[str]] = None
    if args.parallel and args.exec_cmd:
        exec_parallel_cmd = args.exec_cmd

    if_exec_cmd: Optional[list[str]] = 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"{subprocess.list2cmdline(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
                failed_res: CommandResult = CommandResult(
                    command=args.exec_cmd,
                    returncode=err.returncode
                )
                execution_errors.append((repo_path, failed_res))
            except FileNotFoundError:
                failed_res: CommandResult = CommandResult(
                    command=args.exec_cmd,
                    returncode=127
                )
                execution_errors.append((repo_path, failed_res))

        # 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, BrokenPipeError):
        print("\nInterrupting...", file=sys.stderr)
        sys.exit(1)
