import argparse
import json
import re
from typing import Optional

from psi_utils.help_formatter import format_help
from psi_utils.json_utils import JsonEncoder


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: Optional[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)
