import argparse

import yaml


class FunctionTag:
    def __init__(self, value):
        self.value = value


def prompt_func(mode, lang):
    prompt_map = {
        "prompt_1": "Does this statement; {{tweet}} have a Neutral, Positive or Negative sentiment? Labels only",
        "prompt_2": f"Does this {lang} statement; "
        "'{{tweet}}' have a Neutral, Positive or Negative sentiment? Labels only",
        "prompt_3": f"You are an assistant able to detect sentiments in tweets. \n\n"
        f"Given the sentiment labels Neutral, Positive or Negative; what is "
        f"the sentiment of the {lang} statement below? Return only the labels. "
        "\n\ntext: {{tweet}} \nlabel:",
        "prompt_4": "Label the following text as Neutral, Positive, or Negative. Provide only the label as your "
        "response. \n\ntext: {{tweet}} \nlabel: ",
        "prompt_5": f"You are tasked with performing sentiment classification on the following {lang} text. "
        f"For each input, classify the sentiment as positive, negative, or neutral. "
        f"Use the following guidelines: \n\n "
        f"Positive: The text expresses happiness, satisfaction, or optimism. \n"
        f"Negative: The text conveys disappointment, dissatisfaction, or pessimism. \n"
        f"Neutral: The text is factual, objective, or without strong emotional undertones. \n\n"
        f"If the text contains both positive and negative sentiments, choose the dominant sentiment. "
        f"For ambiguous or unclear sentiments, select the label that best reflects the overall tone. "
        "Please provide a single classification for each input.\n\ntext: {{tweet}} \nlabel: ",
    }
    return prompt_map[mode]


def gen_lang_yamls(output_dir: str, overwrite: bool, mode: str) -> None:
    """
    Generate a yaml file for each language.

    :param output_dir: The directory to output the files to.
    :param overwrite: Whether to overwrite files if they already exist.
    """
    err = []
    languages = {
        "amh": "Amharic",
        "arq": "Algerian Arabic",
        "ary": "Moroccan Arabic",
        "hau": "Hausa",
        "ibo": "Igbo",
        "kin": "Kinyarwanda",
        "orm": "Oromo",
        "pcm": "Nigerian Pidgin",
        "por": "Mozambique Portuguese",
        "swa": "Swahili",
        "tir": "Tigrinya",
        "tso": "Xithonga",
        "twi": "Twi",
        "yor": "Yoruba",
    }
    for lang in languages.keys():
        try:
            file_name = f"afrisenti_{lang}.yaml"
            task_name = f"afrisenti_{lang}_{mode}"
            yaml_template = "afrisenti"
            if int(mode.split("_")[-1]) > 1:
                yaml_details = {
                    "include": yaml_template,
                    "task": task_name,
                    "dataset_name": lang,
                    "doc_to_text": prompt_func(mode, languages[lang]),
                }
            else:
                yaml_details = {
                    "include": yaml_template,
                    "task": task_name,
                    "dataset_name": lang,
                }
            with open(
                f"{output_dir}/{mode}/{file_name}",
                "w" if overwrite else "x",
                encoding="utf8",
            ) as f:
                f.write("# Generated by utils.py\n")
                yaml.dump(
                    yaml_details,
                    f,
                    allow_unicode=True,
                )
        except FileExistsError:
            err.append(file_name)

    if len(err) > 0:
        raise FileExistsError(
            "Files were not created because they already exist (use --overwrite flag):"
            f" {', '.join(err)}"
        )


def main() -> None:
    """Parse CLI args and generate language-specific yaml files."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--overwrite",
        default=True,
        action="store_true",
        help="Overwrite files if they already exist",
    )
    parser.add_argument(
        "--output-dir",
        default="./",
        help="Directory to write yaml files to",
    )
    parser.add_argument(
        "--mode",
        default="prompt_1",
        choices=["prompt_1", "prompt_2", "prompt_3", "prompt_4", "prompt_5"],
        help="Prompt number",
    )
    args = parser.parse_args()

    gen_lang_yamls(output_dir=args.output_dir, overwrite=args.overwrite, mode=args.mode)


if __name__ == "__main__":
    main()