import argparse
import json
import re

from psi_utils.help_formatter import format_help


class JsonEncoder(json.JSONEncoder):
    def default(self, obj):
        if not type(obj) in [int, float, complex, dict, tuple, list, bool] and obj is not None:
            return str(obj)

        return json.JSONEncoder.default(self, obj)


class ArgumentParser(argparse.ArgumentParser):
    def __init__(self, default_help_format: bool = False, **kwargs):
        super().__init__(**kwargs)
        self.__default_help_format = bool(default_help_format)

    def parse_args_and_prompt(self):
        args = self.parse_args()
        args_dict = {arg[0]: arg[1] for arg in vars(args).items()}

        user_input = validate_interactive_input(
            prompt=f'Running with args:\n{json.dumps(args_dict, indent=4, cls=JsonEncoder)}\nEnter y/n to continue:',
            allowed_values=['y', 'n']
        )
        return None if user_input != 'y' else args

    def print_help(self, file=None):
        if self.__default_help_format:
            super().print_help(file=file)
        else:
            msg = format_help(self)
            if file:
                file.write(msg)
            else:
                print(msg)


def validate_interactive_input(prompt: str, regex: str = None, allowed_values: list = []) -> str:
    while True:
        user_input = input(f'{prompt}\n').lower().strip()
        if regex and re.match(regex, user_input):
            return user_input
        if allowed_values and user_input in allowed_values:
            return user_input

        msg = '\nInvalid input. '
        msg += f'Input must match regex "{regex}". ' if regex else ''
        msg += f'Input must be one of {allowed_values}.' if allowed_values and not regex else ''
        msg += f'Or input must be one of {allowed_values}.' if allowed_values and regex else ''
        print(msg)
