#!/usr/bin/env python3

from os import path
from collections import defaultdict

from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import assert_equal


# bash cli completion file header
COMPLETION_HEADER = """# Dynamic bash programmable completion for bitcoin-cli(1)
#     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.

"""

# 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_autocomplete(self, pos):
        """ Generate the autocomplete file line relevent 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\" ) )"

# 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(BitcoinTestFramework):
    def set_test_params(self):
        self.num_nodes = 1

    def skip_test_if_missing_module(self):
        self.skip_if_no_cli()
        # self.skip_if_no_wallet()
        self.skip_if_no_bitcoind_zmq()

    def add_options(self, parser):
        parser.add_argument(
                '--header',
                help='Static header part of the bash completion file',
        )

        parser.add_argument(
                '--footer',
                help='Static footer part of the bash completion file',
        )

        parser.add_argument(
                '--completion',
                help='Location of the current bash completion file',
        )

        parser.add_argument(
                '--overwrite',
                default=False,
                action='store_true',
                help='Force the test to overwrite the file pointer to by the --completion'
                     'to the newly generated 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_completion_block(self, options):
        commands = [o.command for o in options]
        self.log.info(f"Generating part of the completion file 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_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_completion_file(self, commands):
        try:
            with open(self.options.header, 'r', encoding='utf-8') as header_file:
                header = header_file.read()

            with open(self.options.footer, 'r', encoding='utf-8') as footer_file:
                footer = footer_file.read()
        except Exception as e:
            raise AssertionError(
                        f"Could not read header/footer ({self.options.header} and {self.options.footer}) files. "
                        f"Tell the test where to find them using the --header/--footer parameters ({e})."
            )
        return COMPLETION_HEADER + header + commands + footer

    def write_completion_file(self, new_file):
        try:
            with open(self.options.completion, '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 {self.options.completion}. "
                        f"Tell the test where to find it using the --completion parameters ({e})."
            )

    def read_completion_file(self):
        try:
            with open(self.options.completion, 'r', encoding='utf-8') as completion_file:
                return completion_file.read()
        except Exception as e:
            raise AssertionError(
                        f"Could not read the autocomplete file ({self.options.completion}) file. "
                        f"Tell the test where to find it using the --completion parameters ({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')
        if self.options.header is None or len(self.options.header) == 0:
            self.options.header = path.join(test_data_dir, 'bitcoin-cli.header.bash-completion')

        if self.options.footer is None or len(self.options.footer) == 0:
            self.options.footer = path.join(test_data_dir, 'bitcoin-cli.footer.bash-completion')

        if self.options.completion is None or len(self.options.completion) == 0:
            self.options.completion = path.join(src_dir, 'contrib', 'completions', 'bash', 'bitcoin-cli.bash')

        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]
        commands = [self.get_command_options(c) for c in commands]

        self.log.info('Generating new autocompletion file...')
        commands = self.generate_completion_block(commands)
        new_completion = self.generate_completion_file(commands)

        if self.options.overwrite:
            self.log.info("Overwriting the completion file...")
            self.write_completion_file(new_completion)

        self.log.info('Checking if the generated and the original completion files matches...')
        completion = self.read_completion_file()
        assert_equal(new_completion, completion)

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