# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import annotations

import json
import subprocess
from collections.abc import Iterable

BASH_COMPLETION_TEMPLATE = """# DO NOT EDIT.
# This script is autogenerated by build-support/bin/generate_completions.py

function pants_completions()
{{
    local -r PANTS_GOALS="\\
        {goals}"

    {goal_options}

    local -r PANTS_GLOBAL_OPTIONS="\\
        {global_options}"

    local current_word previous_word previous_goal

    current_word=${{COMP_WORDS[COMP_CWORD]}}
    previous_word=${{COMP_WORDS[COMP_CWORD-1]}}
    previous_goal=$(get_previous_goal)

    case "$previous_goal" in
    {goal_cases}
    *)
        complete_options="$PANTS_GLOBAL_OPTIONS"
        ;;

    esac

    # Show goals, unless the tab completion has a hyphen - then show global/goal options instead
    if [[ $current_word == -* ]]; then
        COMPREPLY=( $( compgen -W "$complete_options" -- $current_word ))
    else
        COMPREPLY=( $( compgen -W "$PANTS_GOALS" -- $current_word ))
    fi
    return 0
}}

# Get the most recent goal in the command arguments, so options can be correctly applied
# This function will ignore hyphenated options when looking for the goal
# If this is empty, we're at the top-level Pants command
function get_previous_goal()
{{
    local previous_goal i current_word
    previous_goal=
    for (( i=$COMP_CWORD; i > 0; --i )); do
        current_word=${{COMP_WORDS[i]}}
        if [[ $current_word != -* ]]; then
            previous_goal=$current_word
            break
        fi
    done
    echo $previous_goal
}}

complete -o default -o bashdefault -F pants_completions pants
"""

GOAL_OPTIONS_TEMPLATE = """
    local -r PANTS_{name}_OPTIONS="\\
        {options}"
"""

GOAL_CASE_TEMPLATE = """
    {goal})
        complete_options="$PANTS_{bash_name}_OPTIONS"
        ;;
"""


def _run_pants_help():
    process = subprocess.run(["./pants", "help-all"], capture_output=True, check=True)
    process.check_returncode()
    return json.loads(process.stdout)


def _parse_help_for_options(help: dict, scope: str) -> tuple[frozenset[str], frozenset[str]]:
    """scope represents the goal or subsystem of interest Returns a tuple containing the scoped
    options, followed by unscoped."""
    scoped_help_info = help["scope_to_help_info"][scope]

    scoped_options = []
    unscoped_options = []
    for option in scoped_help_info["basic"] + scoped_help_info["advanced"]:
        scoped_options.extend(option["scoped_cmd_line_args"])
        unscoped_options.extend(option["unscoped_cmd_line_args"])
    return (
        frozenset(scoped_options),
        frozenset(unscoped_options),
    )


def _bashify_name(name: str) -> str:
    # Replace hyphens with underscores for variable names
    # Use upper-case goal names
    return name.replace("-", "_").upper()


def _hydrate_bash_template(
    goal_options: dict[str, Iterable[str]], global_options: Iterable[str]
) -> str:
    # Populate the PANTS_GLOBAL_OPTIONS variable
    global_options_str = "\\\n        ".join(sorted(global_options))

    # Populate the PANTS_GOALS variable
    goal_names = sorted(goal_options.keys())
    goals_str = "\\\n        ".join(goal_names)

    # Create variables of the form PANTS_{GOAL}_OPTIONS
    goal_options_templates = ""
    for name, options in sorted(goal_options.items()):
        goal_options_templates += GOAL_OPTIONS_TEMPLATE.format(
            name=_bashify_name(name), options="\\\n        ".join(sorted(options))
        )

    # Enumerate the completion case statements and variable names
    # Note that the no-match goal case is hard-coded as "*" in the template
    goal_cases_templates = ""
    for name in goal_names:
        goal_cases_templates += GOAL_CASE_TEMPLATE.format(goal=name, bash_name=_bashify_name(name))

    return BASH_COMPLETION_TEMPLATE.format(
        global_options=global_options_str,
        goals=goals_str,
        goal_options=goal_options_templates,
        goal_cases=goal_cases_templates,
    )


def main() -> None:
    help_results = _run_pants_help()

    all_scopes: frozenset[str] = frozenset(help_results["scope_to_help_info"].keys())
    goal_scopes: frozenset[str] = frozenset(help_results["name_to_goal_info"].keys())
    pants_scope = ""
    subsystem_scopes = all_scopes.difference([goal_scopes, pants_scope])

    # Holds the scoped options we will complete after "./pants" (i.e. at the global scope)
    # e.g. "./pants -<TAB>" could complete to "./pants --test-use-coverage"
    all_scoped_options: set[str] = set()

    # Holds the unscoped options we will complete after typing a goal
    # e.g. "./pants test -<TAB>" could complete to "./pants test --use-coverage"
    goal_options: dict[str, Iterable[str]] = {}
    for scope in goal_scopes:
        scoped_goal_options, unscoped_goal_options = _parse_help_for_options(help_results, scope)
        goal_options[scope] = unscoped_goal_options
        all_scoped_options.update(scoped_goal_options)

    # Subsystem completion options are applied at the global level
    # e.g. "./pants -<TAB>" could complete to "./pants --pytest-args"
    for name in subsystem_scopes:
        scoped_subsystem_options, _ = _parse_help_for_options(help_results, name)
        all_scoped_options.update(scoped_subsystem_options)

    # Special case for Pants options, the scope name is ""
    # e.g. "./pants -<TAB>" could complete to "./pants --loop"
    scoped_pants_options, _ = _parse_help_for_options(help_results, pants_scope)
    all_scoped_options.update(scoped_pants_options)

    print(_hydrate_bash_template(goal_options, global_options=all_scoped_options))


if __name__ == "__main__":
    main()
