"""Prompt templates for LongLaMP tasks."""


SYSTEM_PROMPT = "You are a helpful writing assistant."


def build_query_prompt(query_input: str, task: str) -> str:
    """Build the inference prompt for a query (no personalization text)."""
    if task == "review":
        return (
            f"{query_input}\n\n"
            f"Write the review text:"
        )
    else:  # topic
        return (
            f"{query_input}\n\n"
            f"Write the post content:"
        )


def build_support_prompt(support_input: str, task: str) -> str:
    """Build the prompt for a support item (used in teacher forcing to cache hidden states)."""
    if task == "review":
        return (
            f"{support_input}\n\n"
            f"Write the review text:"
        )
    else:  # topic
        return (
            f"{support_input}\n\n"
            f"Write the post content:"
        )


def build_prompt_with_examples(query_input: str, support_items: list, task: str) -> str:
    """Build prompt with K support examples included as in-context demonstrations.

    Used by Prompt-All-K and BM25-Top1 baselines.
    """
    parts = []
    parts.append("Here are some examples of this user's writing style:\n")

    for i, item in enumerate(support_items, 1):
        parts.append(f"--- Example {i} ---")
        parts.append(f"Prompt: {item['support_input']}")
        parts.append(f"Response: {item['support_output']}")
        parts.append("")

    parts.append("Now, write in the same style as the examples above.\n")
    parts.append(build_query_prompt(query_input, task))

    return "\n".join(parts)


def build_chat_messages(prompt: str) -> list:
    """Wrap a prompt into chat message format for Qwen."""
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ]