from utils.constants import (
    GPT4O,
    GPT3_5,
    RECURSIVE,
    PARALLEL,
    LATEST_MESSAGE,
    ALL_MESSAGES,
)
import argparse


def map_arguments(arg):
    mapping = {
        "gpt4o": GPT4O,
        "gpt3.5": GPT3_5,
        "r": RECURSIVE,
        "p": PARALLEL,
        "l": LATEST_MESSAGE,
        "a": ALL_MESSAGES,
    }
    return mapping.get(arg, arg)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-l",
        "--llm_type",
        type=map_arguments,
        choices=["4o", "3.5"],
        default=GPT4O,
    )
    parser.add_argument(
        "-i",
        "--infection_mode",
        nargs="+",
        type=map_arguments,
        choices=["r", "p"],
        default=[RECURSIVE, PARALLEL],
    )
    parser.add_argument(
        "-c",
        "--communication_mode",
        nargs="+",
        type=map_arguments,
        choices=["l", "a"],
        default=[LATEST_MESSAGE, ALL_MESSAGES],
    )
    return parser.parse_args()
