"""Prompt renderer plugins: one registered function per ``prompt_style``.

Before this module the runner's ``assemble()`` grew a chain of
``if spec_style == 'xxx'`` branches; now each style is a plugin:

    @register_prompt_renderer('aime_es')
    def aime(question, sample, spec, prompt_style):
        return {'question': ...}

Renderer contract:
- input:  the bare question text + the Sample + the DatasetSpec
- output: dict with any of ``question`` (rewritten), ``system`` (a system
  message to prepend), ``few_shot_header``/``few_shot_glue`` (few-shot
  layout hints consumed by assemble)
- unregistered styles fall back to assemble's generic MCQ/QA handling.

Adding a benchmark prompt style = dropping a plugin here; the runner
never changes.
"""

from typing import Any, Callable, Dict, Optional

RENDERERS: Dict[str, Callable] = {}


def register_prompt_renderer(*styles: str):
    def decorator(fn):
        for s in styles:
            RENDERERS[s] = fn
        return fn

    return decorator


def get_renderer(style: str) -> Optional[Callable]:
    return RENDERERS.get(style)


def render(style: str, question: str, sample, spec, prompt_style: str = '') -> Dict[str, Any]:
    """Apply the style's renderer; unknown styles pass through untouched."""
    fn = RENDERERS.get(style)
    if fn is None:
        return {}
    out = fn(question=question, sample=sample, spec=spec, prompt_style=prompt_style)
    return out if isinstance(out, dict) else {}


# ------------------------------ plugins ------------------------------


@register_prompt_renderer('trivia_es')
def trivia_es(question, sample, spec, prompt_style):
    # es trivia template, verbatim (open-book with wiki evidence, trailing
    # newline included)
    return {'question': (
        'Read the content and answer the following question.\n\n'
        f"Content: {(sample.metadata or {}).get('evidence') or []}\n\n"
        f'Question: {question}\n\n'
        'The last line of your response should be of the form "ANSWER: [ANSWER]" '
        '(without quotes) where [ANSWER] is the answer to the problem.\n')}


@register_prompt_renderer('aime_es')
def aime_es(question, sample, spec, prompt_style):
    # es/MathArena template (NOT lstripped -- leading \n kept; reminder tail
    # after the question, both verbatim from aime_adapter)
    return {'question': (
        '\nSolve the following math problem step by step. '
        'Put your answer inside \\boxed{}.\n\n' + question
        + '\n\nRemember to put your answer inside \\boxed{}.')}


@register_prompt_renderer('imo_es')
def imo_es(question, sample, spec, prompt_style):
    return {'question': (
        f'Problem:\n{question}\n\nPlease reason step by step, and put your '
        f'final answer within \\boxed{{}}.\n')}


@register_prompt_renderer('simple_qa_es')
def simple_qa_es(question, sample, spec, prompt_style):
    return {'question': f'Answer the question:\n\n{question}'}


@register_prompt_renderer('lb2_es')
def lb2_es(question, sample, spec, prompt_style):
    # es longbench-v2 template: <text> wrapper + CoT last-line contract
    letters = 'ABCD'
    opts = '\n'.join(f'{letters[i]}) {c}' for i, c in enumerate(sample.choices or []))
    ctx = (sample.metadata or {}).get('context', '')
    return {'question': (
        'Please read the following text and answer the questions below.\n\n'
        f'<text>\n{ctx}\n</text>\n\n'
        "Answer the following multiple choice question. The last line of your response should be "
        "of the following format: 'ANSWER: [LETTER]' (without quotes) where [LETTER] is one of "
        f'{",".join(letters[:len(sample.choices or [])])}. Think step by step before answering.\n\n'
        f'{question}\n\n{opts}')}


@register_prompt_renderer('lcb_es')
def lcb_es(question, sample, spec, prompt_style):
    # es/official LCB (load_utils + adapter, verbatim): the expert-programmer
    # header is a SYSTEM message (injected by the runner); user keeps
    # ### Question:/### Format:/### Answer
    starter = (sample.metadata or {}).get('starter_code')
    if starter:
        fmt = ('### Format: You will use the following starter code to write the '
               'solution to the problem and enclose your code within delimiters.\n'
               f'```python\n{starter}\n```\n\n')
    else:
        fmt = ('### Format: Read the inputs from stdin solve the problem and write '
               'the answer to stdout (do not directly test on the sample inputs). '
               'Enclose your code within delimiters as follows.\n'
               '```python\n# YOUR CODE HERE\n```\n\n')
    return {'question': (f'### Question:\n{question}\n\n{fmt}### Answer: (use the '
                         'provided format with backticks)\n\n'),
            'system': ('You are an expert Python programmer. You will be given a question '
                       '(problem specification) and will generate a correct Python program '
                       'that matches the specification and passes all tests. You will NOT '
                       'return anything except for the program.')}


@register_prompt_renderer('drop_es')
def drop_es(question, sample, spec, prompt_style):
    # es drop: question block = bare passage + 'Question:' line (es does NOT
    # label the test passage; only exemplars carry labels)
    ps = (sample.metadata or {}).get('passage')
    return {'question': f'{ps}\nQuestion: {question}' if ps else f'Question: {question}'}


@register_prompt_renderer('bbh_es')
def bbh_es(question, sample, spec, prompt_style):
    # es bbh PROMPT_TEMPLATE: the test question is wrapped in the Q:/A:
    # contract (the CoT exemplars already follow this pattern)
    return {'question': (
        'Q: ' + question + '\nA: Let\'s think step by step. Put your final '
        'answer in the format of "So the answer is [ANSWER]" (without quotes '
        'and markdown) where [ANSWER] is the answer to the problem.\n')}


@register_prompt_renderer('cot_letter_plain')
def cot_letter_plain(question, sample, spec, prompt_style):
    # es mmlu-pro USER_PROMPT verbatim: Question:/Options: + 'A) x' --
    # NOTE es renders the TEST question options with PARENS while its
    # exemplars use 'A x' (space); replicate the inconsistency exactly
    if not sample.choices:
        return {}
    letters = 'ABCDEFGHIJ'
    opts = '\n'.join(f'{letters[i]}) {c}' for i, c in enumerate(sample.choices)
                     if i < len(letters))
    return {'question': (
        f'Answer the following multiple choice question. The last line of '
        f"your response should be of the following format: 'ANSWER: [LETTER]' "
        f'(without quotes) where [LETTER] is one of '
        f'{",".join(letters[:len(sample.choices)])}. '
        f'Think step by step before answering.\n\nQuestion:\n{question}\nOptions:\n{opts}\n')}


@register_prompt_renderer('cot_letter_zh')
def cot_letter_zh(question, sample, spec, prompt_style):
    if not sample.choices:
        return {}
    letters = 'ABCDEFGH'
    opts = '\n'.join(f'{letters[i]}) {c}' for i, c in enumerate(sample.choices)
                     if i < len(letters))
    # es cmmlu contract, verbatim (incl. trailing newline)
    return {'question': (
        f'回答下面的单项选择题，请选出其中的正确答案。你的回答的最后一行应该是这样的格式：'
        f'"答案：[LETTER]"（不带引号），其中 [LETTER] 是 {",".join(letters[:len(sample.choices)])} 中的一个。'
        f'请在回答前进行一步步思考。\n\n问题：{question}\n选项：\n{opts}\n')}


@register_prompt_renderer('cot_letter')
def cot_letter(question, sample, spec, prompt_style):
    if not sample.choices:
        return {}
    letters = 'ABCDEFGH'
    opts = '\n'.join(f'{letters[i]}) {c}' for i, c in enumerate(sample.choices)
                     if i < len(letters))
    return {'question': (
        f'Answer the following multiple choice question. The last line of '
        f"your response should be of the following format: 'ANSWER: [LETTER]' "
        f'(without quotes) where [LETTER] is one of {",".join(letters[:len(sample.choices)])}. '
        f'Think step by step before answering.\n\n{question}\n\n{opts}')}