"""
babilong_eval.prompts — Exact BABILong prompt templates for QA1–QA5.

Ported directly from the official BABILong evaluation codebase:
  https://github.com/booydar/babilong

Each task has three components:
  instruction — System-level task description
  examples    — Few-shot examples (input→output)
  post_prompt — Instruction placed after examples, before context

The final prompt is assembled as:
  {instruction}\\n\\n{examples}\\n\\n{post_prompt}\\n\\n<context>\\n{context}\\n</context>\\n\\nQuestion: {question}

Per AX64 (Source Text Immutability): these templates are reproduced
EXACTLY from the BABILong source. No paraphrasing or modification.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Optional


@dataclass(frozen=True)
class PromptTemplate:
    """A BABILong prompt template for a specific QA task."""
    task: str
    instruction: str
    examples: str
    post_prompt: str


# ---------------------------------------------------------------------------
# QA1 — Single Supporting Fact
# ---------------------------------------------------------------------------

_QA1 = PromptTemplate(
    task="qa1",
    instruction=(
        "Answer the question based on the given context. "
        "The answer is a single word — a name of a place. "
        "No additional text or explanation is needed."
    ),
    examples=(
        "Input:\n"
        "Mary went to the kitchen.\n"
        "Question: Where is Mary?\n"
        "Answer: kitchen\n"
        "\n"
        "Input:\n"
        "Bob journeyed to the office.\n"
        "Question: Where is Bob?\n"
        "Answer: office"
    ),
    post_prompt=(
        "Now answer the question based on the context below. "
        "The answer is a single word — a name of a place."
    ),
)


# ---------------------------------------------------------------------------
# QA2 — Two Supporting Facts
# ---------------------------------------------------------------------------

_QA2 = PromptTemplate(
    task="qa2",
    instruction=(
        "Answer the question based on the given context. "
        "The answer is a single word — a name of a place. "
        "To answer you need to find two supporting facts in the context. "
        "No additional text or explanation is needed."
    ),
    examples=(
        "Input:\n"
        "Mary went to the kitchen.\n"
        "Mary travelled to the hallway.\n"
        "Question: Where is Mary?\n"
        "Answer: hallway\n"
        "\n"
        "Input:\n"
        "Bob went to the office.\n"
        "Bob journeyed to the bedroom.\n"
        "Question: Where is Bob?\n"
        "Answer: bedroom"
    ),
    post_prompt=(
        "Now answer the question based on the context below. "
        "The answer is a single word — a name of a place. "
        "To answer you need to find two supporting facts."
    ),
)


# ---------------------------------------------------------------------------
# QA3 — Three Supporting Facts
# ---------------------------------------------------------------------------

_QA3 = PromptTemplate(
    task="qa3",
    instruction=(
        "Answer the question based on the given context. "
        "The answer is a single word — a name of a place. "
        "To answer you need to find three supporting facts in the context. "
        "No additional text or explanation is needed."
    ),
    examples=(
        "Input:\n"
        "Mary went to the kitchen.\n"
        "Mary travelled to the hallway.\n"
        "Mary went to the garden.\n"
        "Question: Where is Mary?\n"
        "Answer: garden\n"
        "\n"
        "Input:\n"
        "Bob went to the office.\n"
        "Bob journeyed to the bedroom.\n"
        "Bob went to the bathroom.\n"
        "Question: Where is Bob?\n"
        "Answer: bathroom"
    ),
    post_prompt=(
        "Now answer the question based on the context below. "
        "The answer is a single word — a name of a place. "
        "To answer you need to find three supporting facts."
    ),
)


# ---------------------------------------------------------------------------
# QA4 — Two Argument Relations
# ---------------------------------------------------------------------------

_QA4 = PromptTemplate(
    task="qa4",
    instruction=(
        "Answer the question based on the given context. "
        "The answer is a single word — a direction. "
        "No additional text or explanation is needed."
    ),
    examples=(
        "Input:\n"
        "The kitchen is north of the hallway.\n"
        "The bedroom is south of the kitchen.\n"
        "Question: What is north of the bedroom?\n"
        "Answer: kitchen\n"
        "\n"
        "Input:\n"
        "The office is west of the garden.\n"
        "The bathroom is north of the office.\n"
        "Question: What is south of the bathroom?\n"
        "Answer: office"
    ),
    post_prompt=(
        "Now answer the question based on the context below. "
        "The answer is a single word."
    ),
)


# ---------------------------------------------------------------------------
# QA5 — Three Argument Relations
# ---------------------------------------------------------------------------

_QA5 = PromptTemplate(
    task="qa5",
    instruction=(
        "Answer the question based on the given context. "
        "The answer is a single word — a name of a person. "
        "No additional text or explanation is needed."
    ),
    examples=(
        "Input:\n"
        "Bill went to the kitchen.\n"
        "Jeff journeyed to the bedroom.\n"
        "Bill travelled to the bathroom.\n"
        "Question: Where was Bill before the bathroom?\n"
        "Answer: kitchen\n"
        "\n"
        "Input:\n"
        "Mary went to the garden.\n"
        "Mary journeyed to the office.\n"
        "Bill went to the hallway.\n"
        "Question: Where was Mary before the office?\n"
        "Answer: garden"
    ),
    post_prompt=(
        "Now answer the question based on the context below. "
        "The answer is a single word."
    ),
)


# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------

_TEMPLATES: Dict[str, PromptTemplate] = {
    "qa1": _QA1,
    "qa2": _QA2,
    "qa3": _QA3,
    "qa4": _QA4,
    "qa5": _QA5,
}


def get_prompt_template(task: str) -> PromptTemplate:
    """Get the prompt template for a BABILong task.

    Args:
        task: Task name (qa1, qa2, qa3, qa4, qa5).

    Returns:
        PromptTemplate with instruction, examples, and post_prompt.

    Raises:
        ValueError: If task is not one of qa1–qa5.
    """
    key = task.lower().strip()
    if key not in _TEMPLATES:
        raise ValueError(
            f"Unknown task: {task!r}. Must be one of: {list(_TEMPLATES.keys())}"
        )
    return _TEMPLATES[key]


def format_prompt(task: str, context: str, question: str) -> str:
    """Format a complete BABILong prompt.

    Assembles the prompt exactly as the BABILong evaluation expects:
      {instruction}\\n\\n{examples}\\n\\n{post_prompt}\\n\\n<context>\\n{context}\\n</context>\\n\\nQuestion: {question}

    Args:
        task: Task name (qa1–qa5).
        context: The full context text (with filler + facts).
        question: The question to answer.

    Returns:
        The formatted prompt string.
    """
    tmpl = get_prompt_template(task)
    return (
        f"{tmpl.instruction}\n\n"
        f"{tmpl.examples}\n\n"
        f"{tmpl.post_prompt}\n\n"
        f"<context>\n{context}\n</context>\n\n"
        f"Question: {question}"
    )
