import argparse
from typing import Callable, Optional, Annotated
from gettext import gettext as _

KommandCompleter = Annotated[
    Callable[[str, Optional[int]], list[str]],
    "Completer for one value of an argument. Takes as input (token, argrank) and returns the list of matches for this token.",
]


class KomandParser(argparse.ArgumentParser):
    """
    Subclass of argparse.ArgumentParser, that implements
    a 'complete' method, which can be called with substrings
    of command line in order to propose matches to complete
    the command line.

    The `exit_on_error` parameter of ArgumentParser is forced to
    False, and the `prog` parameter is made mandatory to enforce
    independance with sys.argv.

    Matches are computed by `completer` callables given to
    each argument on `argparse.ArgumentParser.add_argument`
    call.

    Those methods must take as input the current token which has
    to be completed, as well as its rank in list of arguments.
    """

    def __init__(self, prog: str, *args, **kwargs):
        self.dest_completers: dict[str, KommandCompleter] = {}
        self.all_option_strings: list[str] = []
        super().__init__(prog=prog, *args, **kwargs, exit_on_error=False)

    def add_argument(
        self,
        *args,
        completer: KommandCompleter | None = None,
        **kwargs,
    ) -> argparse.Action:
        """
        Calls super().add_argument. To add a custom callable that
        completes the argument values, add a parameter `completer`
        when calling this method.
        The signature of the completer is : [str, int], list[str].
        """
        action = super().add_argument(*args, **kwargs)
        self.all_option_strings += action.option_strings

        def empty_fill(x, rank=None):
            return [x]

        if completer is None:
            completer = empty_fill
        self.dest_completers[action.dest] = completer
        return action

    @staticmethod
    def _find_matching_words(word: str, word_list: list[str]) -> list[str]:
        """
        Return the subset of word_list which starts exactly as
        word.

        Parameters :
        ---
            - word (str): the beginning of a word
            - word_list (list[str]) : the list of strings in which we
                want to find words that starts as `word`

        Returns :
        ---
            The list of words which starts exactly as `word`. Possibly
            an empty list.
        """
        all_matches = []
        for each_word in word_list:
            if len(each_word) < len(word):
                continue
            potential_match = each_word[: len(word)]
            if potential_match == word:
                all_matches.append(each_word)
        return all_matches

    def _classify_tokens(self, arg_strings: list[str]):
        """
        Method extracted from super()._parse_known_args
        Used to classify if each token corresponds to
        optional, positional argument or is '--'.

        Parameters :
        ---
            - arg_strings (list[str]) : the list of tokens

        Returns :
        ---
            A tuple (arg_string_pattern_parts, option_string_indices).
            First element is :
                A list of same length as arg_strings; where values are :
                    - 'A' if token is for positional argument
                    - 'O' if token is for optional argument
                    - '-' if token is '--'
            Second is :
                A dict where keys are those of optional arguments, and values
                are the Action object corresponding to
        """
        option_string_indices = {}
        arg_string_pattern_parts = []
        arg_strings_iter = iter(arg_strings)
        for i, arg_string in enumerate(arg_strings_iter):

            # all args after -- are non-options
            if arg_string == "--":
                arg_string_pattern_parts.append("-")
                for arg_string in arg_strings_iter:
                    arg_string_pattern_parts.append("A")

            # otherwise, add the arg to the arg strings
            # and note the index if it was an option
            else:
                option_tuples = self._parse_optional(arg_string)
                if option_tuples is None:
                    pattern = "A"
                else:
                    option_string_indices[i] = option_tuples
                    pattern = "O"
                arg_string_pattern_parts.append(pattern)

        return arg_string_pattern_parts, option_string_indices

    def complete(self, tokens: list[str]) -> list[str]:
        """
        Completes the last token of the list, by acting
        similarly as argparse.parse_args to call the
        proper `complete` method of arguments.

        If the last token is an option string (--option_name for
        example), returns the list of option names that are not
        yet written in the command line.

        Parameters :
        ---
            - tokens (list[str]) : the list of tokens, from
                which the last token will be completed.

        Returns :
        ---
            The list of matching words for the last token.
        """
        try:
            return self._complete(tokens)
        except argparse.ArgumentError:
            if len(tokens) > 0:
                return [tokens[-1]]
            else:
                return [""]

    def _complete(self, tokens: list[str]) -> list[str]:
        """
        Completes the command line by calling 'completer' method
        of the last argument, or by giving appropriate option names.

        Parameters:
        ---
            - tokens (list[str]) : the list of tokens

        Returns :
        ---
            - the list of matches for the last token
        """

        token_kinds, option_string_indices = self._classify_tokens(tokens)

        if len(tokens) == 0:  # if no token is provided
            positionals = self._get_positional_actions()
            if len(positionals) > 0:
                first_pos = positionals[0]
                completer = self.dest_completers[first_pos.dest]
                return completer("", 1)

            if len(self.all_option_strings) > 0:
                return self.all_option_strings
            return [""]

        last_token_is_not_empty = len(tokens[-1]) > 0
        if last_token_is_not_empty and tokens[-1][0] == "-":
            # then we complete with all the name of the optional that
            # is not yet in the command line, and that matches the beginning
            # of the token
            if "=" not in tokens[-1]:
                # if '=', then we have an explicit arg and we let the code
                # below do its job

                unused_option_strings = []
                for each_action in self._actions:
                    is_action_in_list = False
                    for each_option_string in each_action.option_strings:
                        if each_option_string in tokens:
                            is_action_in_list = True
                    if not is_action_in_list:
                        unused_option_strings += each_action.option_strings

                return self._find_matching_words(tokens[-1], unused_option_strings)

        def consume_optional(start_index):
            """
            Directly extracted from super()._parse_known_args
            Small modifications to deal with the completion need :
                - no warning for deprecation,
                - the action is not triggered (no take_action)
            """
            # get the optional identified at this index
            option_tuples = option_string_indices[start_index]

            action, option_string, sep, explicit_arg = option_tuples[0]

            # identify additional optionals in the same arg string
            # (e.g. -xyz is the same as -x -y -z if no args are required)
            match_argument = self._match_argument
            action_tuples = []
            while True:
                # if there is an explicit argument, try to match the
                # optional's string arguments to only this
                if explicit_arg is not None:
                    arg_count = match_argument(action, "A")

                    # if the action is a single-dash option and takes no
                    # arguments, try to parse more single-dash options out
                    # of the tail of the option string
                    chars = self.prefix_chars
                    if (
                        arg_count == 0
                        and option_string[1] not in chars
                        and explicit_arg != ""
                    ):
                        if sep or explicit_arg[0] in chars:
                            msg = _("ignored explicit argument %r")
                            raise argparse.ArgumentError(action, msg % explicit_arg)
                        action_tuples.append((action, [], option_string))
                        char = option_string[0]
                        option_string = char + explicit_arg[0]
                        optionals_map = self._option_string_actions
                        if option_string in optionals_map:
                            action = optionals_map[option_string]
                            explicit_arg = explicit_arg[1:]
                            if not explicit_arg:
                                sep = explicit_arg = None
                            elif explicit_arg[0] == "=":
                                sep = "="
                                explicit_arg = explicit_arg[1:]
                            else:
                                sep = ""
                        else:
                            extras.append(char + explicit_arg)
                            extras_pattern.append("O")
                            stop = start_index + 1
                            break
                    # if the action expect exactly one argument, we've
                    # successfully matched the option; exit the loop
                    elif arg_count == 1:
                        stop = start_index + 1
                        args = [explicit_arg]
                        action_tuples.append((action, args, option_string))
                        break

                    # error if a double-dash option did not use the
                    # explicit argument
                    else:
                        msg = _("ignored explicit argument %r")
                        raise argparse.ArgumentError(action, msg % explicit_arg)

                # if there is no explicit argument, try to match the
                # optional's string arguments with the following strings
                # if successful, exit the loop
                else:
                    start = start_index + 1
                    selected_patterns = arg_strings_pattern[start:]
                    arg_count = match_argument(action, selected_patterns)
                    stop = start + arg_count
                    args = tokens[start:stop]
                    action_tuples.append((action, args, option_string))
                    break

            return stop

        # will be changed by consuming positials after
        positionals = self._get_positional_actions()
        arg_strings_pattern = "".join(token_kinds)

        def consume_positionals(start_index):
            """
            Directly extracted from super()._parse_known_args
            Small modifications to deal with the completion need :
                - no warning for deprecation,
                - the action is not triggered (no take_action)
            """
            # match as many Positionals as possible
            match_partial = self._match_arguments_partial
            selected_pattern = arg_strings_pattern[start_index:]
            arg_counts = match_partial(positionals, selected_pattern)

            # slice off the appropriate arg strings for each Positional
            # and add the Positional and its args to the list
            for action, arg_count in zip(positionals, arg_counts):
                args = tokens[start_index : start_index + arg_count]
                # Strip out the first '--' if it is not in REMAINDER arg.
                if action.nargs == argparse.PARSER:
                    if arg_strings_pattern[start_index] == "-":
                        assert args[0] == "--"
                        args.remove("--")
                elif action.nargs != argparse.REMAINDER:
                    if (
                        arg_strings_pattern.find(
                            "-", start_index, start_index + arg_count
                        )
                        >= 0
                    ):
                        args.remove("--")
                start_index += arg_count

            # slice off the Positionals that we just parsed and return the
            # index at which the Positionals' string args stopped
            positionals[:] = positionals[len(arg_counts) :]
            return start_index

        # again : copy-pasted from super()._parse_known_args
        # add a trigger that detect when we are at the last token
        extras = []
        extras_pattern = []
        start_index = 0
        last_action = None
        if option_string_indices:
            max_option_string_index = max(option_string_indices)
        else:
            max_option_string_index = -1
            arg_counter = 0

            for positional in positionals:
                if positional.nargs is None:
                    nargs = 1
                elif positional.nargs == "*":
                    nargs = 1  # we still count it even if it is
                    # optional
                elif isinstance(positional.nargs, str):
                    try:
                        nargs = int(positional.nargs)
                    except Exception as e:
                        nargs = 0
                else:
                    nargs = positional.nargs

                arg_counter += nargs

                if arg_counter >= len(tokens):
                    rank = len(tokens) - arg_counter + nargs
                    last_action = positional
                    completer = self.dest_completers[last_action.dest]
                    return completer(tokens[-1], rank)

        arg_rank = 1
        while start_index <= max_option_string_index:

            # find the next_option_string_index :
            next_option_string_index = start_index
            while next_option_string_index <= max_option_string_index:
                if next_option_string_index in option_string_indices:
                    break
                next_option_string_index += 1

            # consume any Positionals preceding the next option
            if start_index != next_option_string_index:

                positionals_end_index = consume_positionals(start_index)

                # only try to parse the next optional if we didn't consume
                # the option string during the positionals parsing
                if positionals_end_index > start_index:
                    start_index = positionals_end_index
                    continue
                else:
                    start_index = positionals_end_index

            # consume the next optional and any arguments for it

            try:
                new_start_index = consume_optional(start_index)
            except argparse.ArgumentError as e:
                last_action = option_string_indices[start_index][0][0]
                arg_rank = len(tokens) - 1 - start_index
                break
            else:
                if new_start_index >= len(tokens):
                    last_action = option_string_indices[start_index][0][0]
                    arg_rank = new_start_index - 1 - start_index
                    if arg_rank == 0:
                        arg_rank = 1
                    break
                else:
                    last_action = positionals[0] if len(positionals) > 0 else None
                    start_index = new_start_index
                    arg_rank = len(tokens) - start_index

        if last_action is None:
            return [tokens[-1]]

        completer = self.dest_completers[last_action.dest]
        return completer(tokens[-1], arg_rank)


if __name__ == "__main__":
    # Example Usage
    import shlex

    parser = KomandParser(prog="my-parser")

    def pos1_completer(word, rank):
        return [f"pos1 completer : {word}-{rank}"]

    def pos2_completer(word, rank):
        return [f"pos2 completer : {word}-{rank}"]

    def opt1_completer(word, rank):
        return [f"opt1 completer : {word}-{rank}"]

    def opt2_completer(word, rank):
        return [f"opt2 completer : {word}-{rank}"]

    parser.add_argument("pos1", nargs=1, completer=pos1_completer)
    parser.add_argument("pos2", nargs=1, completer=pos2_completer)
    parser.add_argument("--opt1", "-o1", nargs=2, completer=opt1_completer)
    parser.add_argument("--opt2", "-o2", nargs=2, completer=opt2_completer)

    arg_strings = shlex.split("argpos1 argpo")
    print(parser.complete(arg_strings))  # ['pos2 completer : argpo-1']

    arg_strings = shlex.split("argpos")
    print(parser.complete(arg_strings))  # ['pos1 completer : argpos-1']

    arg_strings = shlex.split("--")
    print(
        parser.complete(arg_strings)
    )  # ['--help', '--opt1', '--opt2'] - help flag is by default on all argparse.ArgumentParser

    arg_strings = shlex.split("--op")
    print(
        parser.complete(arg_strings)
    )  # ['--opt1', '--opt2'] - help flag is by default on all argparse.ArgumentParser

    arg_strings = shlex.split("argpos1 --opt1 argopt1 arg")
    print(
        parser.complete(arg_strings)
    )  # ['opt1 completer : arg-2'] - second argument of opt1 needs to be completed

    arg_strings = shlex.split("argpos1 --opt1 argopt1 argopt2 arg")
    print(
        parser.complete(arg_strings)
    )  # ['pos2 completer : arg-1'] - completer of pos2

    arg_strings = shlex.split("argpos1 --opt1 argopt1 argopt2 --o")
    print(
        parser.complete(arg_strings)
    )  # ['--opt2'] - completes flag name among the ones not already seen before

    arg_strings = [""]
    print(
        parser.complete(arg_strings)
    )  # ['pos1 completer : -1'] : an empty strings works as well

    arg_strings = shlex.split("argpos1") + [""]
    print(
        parser.complete(arg_strings)
    )  # ['pos2 completer : -1'] : empty string completes here the second positional since the
    # the first has already been filled

    parser_2 = KomandParser(prog="my-second-parser")
    parser_2.add_argument("pos1", nargs="*", completer=pos1_completer)
    parser_2.add_argument("-l", completer=opt1_completer)

    arg_strings = [""]
    print(
        parser_2.complete(arg_strings)
    )  # ['pos1 completer : -1'] : optional positional is supported, only for * (? is unimplementable)

    # expected :
    # ['pos2 completer : argpo-1']
    # ['pos1 completer : argpos-1']
    # ['--help', '--opt1', '--opt2']
    # ['--opt1', '--opt2']
    # ['opt1 completer : arg-2']
    # ['pos2 completer : arg-1']
    # ['--opt2']
    # ['pos1 completer : -1']
    # ['pos2 completer : -1']
    # ['pos1 completer : -1']
