#!/usr/bin/env python3

from os import path
from collections import defaultdict

from test_framework.test_framework import LimenkaTestFramework
from test_framework.util import assert_equal


# Common warning for auto-generated completion files
COMPLETION_WARNING = """#     DO NOT EDIT THIS FILE BY HAND -- THIS WILL FAIL THE FUNCTIONAL TEST tool_cli_completion
# This file is auto-generated by the functional test tool_cli_completion.
# If you want to modify this file, modify test/functional/tool_cli_completion.py and re-autogenerate
# this file via the --overwrite test flag.

"""

# Completion file headers for different shells
BASH_COMPLETION_HEADER = f"""# Dynamic bash programmable completion for limenka-cli(1)
{COMPLETION_WARNING}"""

ZSH_COMPLETION_HEADER = f"""#compdef limenka-cli
# zsh completion for limenka-cli(1)
{COMPLETION_WARNING}"""

# option types which are limited to certain values
TYPED_OPTIONS = [
        ["estimate_mode", {"UNSET", "ECONOMICAL", "CONSERVATIVE"}],
        ["sighashtype", {"ALL", "NONE", "SINGLE", "ALL|ANYONECANPAY",
                         "NONE|ANYONECANPAY", "SINGLE|ANYONECANPAY"}]
]


class PossibleArgs():
    """ Helper class to store options associated to a command. """
    def __init__(self, command):
        self.command = command
        self.arguments = {}

    def set_args(self, position, values):
        """ Set the position-th positional argument as having values as possible values. """
        if position in self.arguments:
            raise AssertionError(f"The positional parameter at position {position} is already defined for command '{self.command}'")

        self.arguments[position] = values
        return self

    def set_bool_args(self, position):
        return self.set_args(position, {"true", "false"})

    def set_file_args(self, position):
        # We consider an empty string as a file value for the sake of simplicity (don't
        # have to create an extra level of indirection).
        return self.set_args(position, {""})

    def set_unknown_args(self, position):
        return self.set_args(position, {})

    def set_typed_option(self, position, arg_name):
        """ Checks if arg_name is a typed option; if it is, sets it and return True. """
        for option_type in TYPED_OPTIONS:
            if arg_name == option_type[0]:
                self.set_args(position, option_type[1])
                return True
        return False

    def has_option(self, position):
        return position in self.arguments and len(self.arguments[position]) > 0

    def get_num_args(self):
        """ Return the max number of positional argument the option accepts. """
        pos = list(self.arguments.keys())
        if len(pos) == 0:
            return 0

        return max(pos)

    def generate_bash_autocomplete(self, pos):
        """ Generate the bash autocomplete file line relevant to the given position pos. """
        if len(self.arguments[pos]) == 0:
            raise AssertionError(f"generating undefined arg id {pos} ({self.arguments})")

        # handle special file case
        if len(self.arguments[pos]) == 1 and len(next(iter(self.arguments[pos]))) == 0:
            return "_filedir"

        # a set order is undefined, so we order args alphabetically
        args = list(self.arguments[pos])
        args.sort()
        return "COMPREPLY=( $( compgen -W \"" + ' '.join(args) + "\" -- \"$cur\" ) )"

    def generate_zsh_autocomplete(self, pos):
        """ Generate the zsh autocomplete file line relevant to the given position pos. """
        if len(self.arguments[pos]) == 0:
            raise AssertionError(f"generating undefined arg id {pos} ({self.arguments})")

        # handle special file case
        if len(self.arguments[pos]) == 1 and len(next(iter(self.arguments[pos]))) == 0:
            return "_files"

        # a set order is undefined, so we order args alphabetically
        args = list(self.arguments[pos])
        args.sort()
        return "_values 'arg' " + ' '.join(f"'{arg}'" for arg in args)

# commands where the option type can only be difficultly derived from the help message
SPECIAL_OPTIONS = [
        PossibleArgs("addnode").set_args(2, {"add", "remove", "onetry"}),
        PossibleArgs("setban").set_args(2, {"add", "remove"}),
]


def generate_start_complete(cword):
    """ Generate the start of an autocomplete block (beware of indentation). """
    if cword > 1:
        return f"""    if ((cword > {cword})); then
        case ${{words[cword-{cword}]}} in"""

    return "    case \"$prev\" in"


def generate_end_complete(cword):
    """ Generate the end of an autocomplete block. """
    if cword > 1:
        return f"\n{' ' * 8}esac\n{' ' * 4}fi\n\n"

    return f"\n{' ' * 4}esac\n"


class CliCompletionTest(LimenkaTestFramework):
    def set_test_params(self):
        self.num_nodes = 1

    def skip_test_if_missing_module(self):
        self.skip_if_no_cli()

    def add_options(self, parser):
        parser.add_argument(
                '--overwrite',
                default=False,
                action='store_true',
                help='Force the test to overwrite the completion files with newly generated ones',
        )
        parser.add_argument(
                '--bash-completion',
                default=None,
                help='Location of the current bash completion file',
        )
        parser.add_argument(
                '--zsh-completion',
                default=None,
                help='Location of the current zsh completion file',
        )
    def parse_single_helper(self, option):
        """ Complete the arguments of option via the RPC format command. """

        res = self.nodes[0].format(command=option.command, output='args_cli')
        if len(res) == 0:
            return option

        if res.count('\n') > 1:
            raise AssertionError(
                f"command {option.command} doesn't support format RPC. Should it be a hidden command? "
                f"Please call RPCHelpMan::Check when adding a new non-hidden command. Returned: {res}"
            )

        for idx, argument in enumerate(res.split(",")):
            elems = argument.split(":")

            if option.set_typed_option(idx+1, elems[0]):
                continue

            if elems[1] == "boolean":
                option.set_bool_args(idx+1)
                continue

            if elems[1] == "file":
                option.set_file_args(idx+1)
                continue

            if not option.has_option(idx+1):
                option.set_unknown_args(idx+1)

        return option

    def get_command_options(self, command):
        """ Returns the corresponding PossibleArgs for the command. """

        # verify it's not a special option first
        for soption in SPECIAL_OPTIONS:
            if command == soption.command:
                return self.parse_single_helper(soption)

        return self.parse_single_helper(PossibleArgs(command))

    def generate_bash_completion_block(self, options):
        """Generate bash-specific completion block."""
        commands = [o.command for o in options]
        self.log.info(f"Generating bash completion for options {commands}")

        if len(options) == 0:
            return ""

        generated = ""
        max_pos_options = max(options, key=lambda o: o.get_num_args()).get_num_args()
        for cword in range(max_pos_options, 0, -1):
            this_options = [option for option in options if option.has_option(cword)]
            if len(this_options) == 0:
                continue

            # group options by their arguments value
            grouped_options = defaultdict(list)
            for option in this_options:
                arg = option.generate_bash_autocomplete(cword)
                grouped_options[arg].append(option)

            # generate the cword block
            indent = 12 if cword > 1 else 8
            generated += generate_start_complete(cword)
            for line, opt_gr in grouped_options.items():
                opt_gr.sort(key=lambda o: o.command)  # show options alphabetically for clarity
                args = '|'.join([o.command for o in opt_gr])
                generated += f"\n{' '*indent}{args})\n"
                generated += f"{' ' * (indent + 4)}{line}\n{' ' * (indent + 4)}return 0\n{' ' * (indent + 4)};;"
            generated += generate_end_complete(cword)

        return generated

    def generate_zsh_completion_block(self, options):
        """Generate zsh-specific completion block."""
        commands = [o.command for o in options]
        self.log.info(f"Generating zsh completion for options {commands}")

        if len(options) == 0:
            return ""

        generated = ""
        max_pos_options = max(options, key=lambda o: o.get_num_args()).get_num_args()

        # Generate completion blocks from highest position to lowest
        for cword in range(max_pos_options, 0, -1):
            this_options = [option for option in options if option.has_option(cword)]
            if len(this_options) == 0:
                continue

            # Group options by their arguments value
            grouped_options = defaultdict(list)
            for option in this_options:
                arg = option.generate_zsh_autocomplete(cword)
                grouped_options[arg].append(option)

            # Generate the CURRENT check and case block
            if cword > 1:
                generated += f"\n    if (( CURRENT > {cword + 1} )); then\n"
                generated += f"        case ${{words[CURRENT-{cword}]}} in\n"
                indent = 12
            else:
                generated += "\n    # Handle previous word completions\n"
                generated += '    case "${words[CURRENT-1]}" in\n'
                indent = 8

            for line, opt_gr in grouped_options.items():
                opt_gr.sort(key=lambda o: o.command)  # show options alphabetically for clarity
                args = '|'.join([o.command for o in opt_gr])
                generated += f"{' '*indent}{args})\n"
                generated += f"{' ' * (indent + 4)}{line}\n"
                generated += f"{' ' * (indent + 4)}return 0\n"
                generated += f"{' ' * (indent + 4)};;\n"

            if cword > 1:
                generated += "        esac\n"
                generated += "    fi\n"
            else:
                generated += "    esac\n"

        return generated

    def generate_both_completion_blocks(self, options):
        """Generate both bash and zsh completion blocks."""
        bash_block = self.generate_bash_completion_block(options)
        zsh_block = self.generate_zsh_completion_block(options)
        return bash_block, zsh_block

    def generate_completion_files(self, bash_commands, zsh_commands, bash_header_path, bash_footer_path, zsh_header_path, zsh_footer_path):
        """Generate both bash and zsh completion files."""
        # Read bash header and footer
        try:
            with open(bash_header_path, 'r', encoding='utf-8') as f:
                bash_header = f.read()
            with open(bash_footer_path, 'r', encoding='utf-8') as f:
                bash_footer = f.read()
        except Exception as e:
            raise AssertionError(
                f"Could not read bash header/footer files ({bash_header_path} and {bash_footer_path}): {e}"
            )

        # Read zsh header and footer
        try:
            with open(zsh_header_path, 'r', encoding='utf-8') as f:
                zsh_header = f.read()
            with open(zsh_footer_path, 'r', encoding='utf-8') as f:
                zsh_footer = f.read()
        except Exception as e:
            raise AssertionError(
                f"Could not read zsh header/footer files ({zsh_header_path} and {zsh_footer_path}): {e}"
            )

        bash_completion = BASH_COMPLETION_HEADER + bash_header + bash_commands + bash_footer
        zsh_completion = ZSH_COMPLETION_HEADER + zsh_header + zsh_commands + zsh_footer

        return bash_completion, zsh_completion

    def write_completion_file(self, new_file, file_path):
        """Write a completion file to the specified path."""
        try:
            with open(file_path, 'w', encoding='utf-8') as completion_file:
                completion_file.write(new_file)
        except Exception as e:
            raise AssertionError(
                f"Could not write the autocomplete file to {file_path}: {e}"
            )

    def read_completion_file(self, file_path):
        """Read a completion file from the specified path."""
        try:
            with open(file_path, 'r', encoding='utf-8') as completion_file:
                return completion_file.read()
        except Exception as e:
            raise AssertionError(
                f"Could not read the autocomplete file ({file_path}): {e}"
            )


    def run_test(self):
        # self.config is not available in self.add_options, so complete filepaths here
        src_dir = self.config["environment"]["SRCDIR"]
        test_data_dir = path.join(src_dir, 'test', 'functional', 'data', 'completion')

        # Define all file paths
        bash_header_path = path.join(test_data_dir, 'limenka-cli.header.bash-completion')
        bash_footer_path = path.join(test_data_dir, 'limenka-cli.footer.bash-completion')

        # Use command line parameter if provided, otherwise use default path
        if self.options.bash_completion:
            bash_completion_path = self.options.bash_completion
        else:
            bash_completion_path = path.join(src_dir, 'contrib', 'completions', 'bash', 'limenka-cli.bash')

        zsh_header_path = path.join(test_data_dir, 'limenka-cli.header.zsh-completion')
        zsh_footer_path = path.join(test_data_dir, 'limenka-cli.footer.zsh-completion')

        # Use command line parameter if provided, otherwise use default path
        if self.options.zsh_completion:
            zsh_completion_path = self.options.zsh_completion
        else:
            zsh_completion_path = path.join(src_dir, 'contrib', 'completions', 'zsh', 'limenka-cli.zsh')

        self.log.info('Parsing help commands to get all the command arguments...')
        commands = self.nodes[0].help().split("\n")
        commands = [c.split(' ')[0] for c in commands if not c.startswith("== ") and len(c) > 0]
        command_options = [self.get_command_options(c) for c in commands]

        self.log.info('Generating new bash and zsh completion files...')
        bash_commands, zsh_commands = self.generate_both_completion_blocks(command_options)

        bash_completion, zsh_completion = self.generate_completion_files(
            bash_commands, zsh_commands,
            bash_header_path, bash_footer_path,
            zsh_header_path, zsh_footer_path
        )

        if self.options.overwrite:
            self.log.info("Overwriting the bash and zsh completion files...")
            self.write_completion_file(bash_completion, bash_completion_path)
            self.write_completion_file(zsh_completion, zsh_completion_path)

        # Check bash completion file
        if path.exists(bash_completion_path):
            self.log.info('Checking if the generated and original bash completion files match...')
            existing_bash = self.read_completion_file(bash_completion_path)
            assert_equal(bash_completion, existing_bash)
        else:
            self.log.warning(f'Bash completion file not found at {bash_completion_path}, skipping comparison')

        # Check zsh completion file
        if path.exists(zsh_completion_path):
            self.log.info('Checking if the generated and original zsh completion files match...')
            existing_zsh = self.read_completion_file(zsh_completion_path)
            assert_equal(zsh_completion, existing_zsh)
        else:
            self.log.warning(f'Zsh completion file not found at {zsh_completion_path}, skipping comparison')

if __name__ == '__main__':
    CliCompletionTest(__file__).main()