from typing import List, Union, Dict
import re, ast
import dataclasses
import tyro
from typing import Type, Tuple, TypeVar


def parse_unknown_args(args: List[str]) -> Dict[str, Union[str, bool, int, float, List]]:
    """
    Parse unknown arguments from the command line.
    """

    arg_name_match = re.compile(r"^--?([a-zA-Z_][a-zA-Z0-9_\.]*)$")
    parsed_args: Dict[str, Union[str, bool, int, float, List[Union[str, bool, int, float]]]] = {}
    arg_names: List[str] = []
    arg_values: List[List[str]] = []

    def literal_eval_str(str_value: str) -> Union[str, int, float, bool]:
        """
        Safely evaluate a string to its literal value.
        """
        try:
            return ast.literal_eval(str_value)
        except (ValueError, SyntaxError):
            return str_value

    for arg in args:
        if arg_name_match.match(arg):
            # This is a valid argument name
            arg_name, *_ = arg_name_match.findall(arg)
            if arg_name not in parsed_args:
                temp_list = []
                parsed_args[arg_name] = temp_list
                arg_names.append(arg_name)
                arg_values.append(temp_list)
        else:
            arg_values[-1].append(arg)

    # Convert lists to appropriate types
    for i, arg_name in enumerate(arg_names):
        if len(arg_values[i]) == 0:
            parsed_args[arg_name] = True
        else:
            values = list(map(literal_eval_str, arg_values[i]))
            if len(values) == 1:
                parsed_args[arg_name] = values[0]
            else:
                parsed_args[arg_name] = values

    # print(f"Parsed unknown args: {parsed_args}")
    # input()
    return parsed_args


T = TypeVar("T")


def auto_extract_args(ArgClass: Type[T]) -> Tuple[T, Dict[str, Union[str, bool, int, float, List]]]:
    """
    Automatically extract arguments from the command line.
    This function is a wrapper around parse_unknown_args to provide a more user-friendly interface.
    """

    parsed_args, args = tyro.cli(
        ArgClass,
        return_unknown_args=True,
    )

    unknown_args = parse_unknown_args(args)

    return parsed_args, unknown_args
