# Historical source snapshot. Comments about accuracy, data origin, or compliance
# are author claims, not verification; see training/README.md and REPRODUCIBILITY.md.
"""
SinhalaMMLU Official Prompt Templates
Derived from the EMNLP 2025 benchmark paper:
'SinhalaMMLU: A Comprehensive Benchmark for Evaluating Multitask Language Understanding in Sinhala'
(Pramodya et al., 2025)
"""

SINHALA_PROMPT_4OPT = """මෙය {subject} විෂයට අදාළ බහුවරණ ප්‍රශ්නයකි. පහත ප්‍රශ්නයට 1, 2, 3, 4 යන පිළිතුරු වලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරන්න.
ප්‍රශ්නය: {question}
1. {opt1}
2. {opt2}
3. {opt3}
4. {opt4}
පිළිතුර: """

SINHALA_PROMPT_5OPT = """මෙය {subject} විෂයට අදාළ බහුවරණ ප්‍රශ්නයකි. පහත ප්‍රශ්නයට 1, 2, 3, 4, 5 යන පිළිතුරු වලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරන්න.
ප්‍රශ්නය: {question}
1. {opt1}
2. {opt2}
3. {opt3}
4. {opt4}
5. {opt5}
පිළිතුර: """


def build_prompt(question_dict: dict) -> str:
    """
    Builds the standardized Sinhala prompt based on question dictionary.
    Handles 'choices' (list) and 'options' (dict or list).
    """
    subject = question_dict.get("subject", "").strip()
    question = question_dict.get("question", "").strip()

    # The dataset uses 'choices', but support 'options' as well
    choices = question_dict.get("choices")
    if choices is None:
        choices = question_dict.get("options", [])

    if isinstance(choices, dict):
        keys = sorted(choices.keys(), key=lambda x: str(x))
        opt_list = [choices[k] for k in keys]
    elif isinstance(choices, list):
        opt_list = choices
    else:
        opt_list = []

    if len(opt_list) >= 5:
        return SINHALA_PROMPT_5OPT.format(
            subject=subject or "සාමාන්‍ය",
            question=question,
            opt1=opt_list[0],
            opt2=opt_list[1],
            opt3=opt_list[2],
            opt4=opt_list[3],
            opt5=opt_list[4],
        )
    elif len(opt_list) >= 4:
        return SINHALA_PROMPT_4OPT.format(
            subject=subject or "සාමාන්‍ය",
            question=question,
            opt1=opt_list[0],
            opt2=opt_list[1],
            opt3=opt_list[2],
            opt4=opt_list[3],
        )
    else:
        opts_formatted = "\n".join([f"{i+1}. {opt}" for i, opt in enumerate(opt_list)])
        return f"ප්‍රශ්නය: {question}\n{opts_formatted}\nපිළිතුර: "