"""
三种模式的 messages 组装逻辑。

Few-shot CoT 的 messages 结构（重点）：
  system
  + 示例1 user + 示例1 assistant
  + 示例2 user + 示例2 assistant
  + 新题 user
"""

from __future__ import annotations

from examples import FEW_SHOT_EXAMPLES, FewShotExample
from llm_client import call_llm
from prompts import (
    SYSTEM_DIRECT,
    SYSTEM_FEW_SHOT_COT,
    SYSTEM_ZERO_SHOT_COT,
    USER_DIRECT_TEMPLATE,
    USER_FEW_SHOT_COT_TEMPLATE,
    USER_ZERO_SHOT_COT_TEMPLATE,
)
from question import TargetQuestion


def _append_few_shot_pairs(
    messages: list[dict[str, str]],
    examples: list[FewShotExample],
) -> None:
    for example in examples:
        messages.append({"role": "user", "content": example.user})
        messages.append({"role": "assistant", "content": example.assistant})


def run_direct(question: TargetQuestion) -> str:
    messages = [
        {"role": "system", "content": SYSTEM_DIRECT},
        {"role": "user", "content": USER_DIRECT_TEMPLATE.format(question=question.question)},
    ]
    return call_llm(messages)


def run_zero_shot_cot(question: TargetQuestion) -> str:
    messages = [
        {"role": "system", "content": SYSTEM_ZERO_SHOT_COT},
        {
            "role": "user",
            "content": USER_ZERO_SHOT_COT_TEMPLATE.format(question=question.question),
        },
    ]
    return call_llm(messages)


def run_few_shot_cot(
    question: TargetQuestion,
    examples: list[FewShotExample] | None = None,
) -> str:
    shot_examples = examples if examples is not None else FEW_SHOT_EXAMPLES
    messages: list[dict[str, str]] = [{"role": "system", "content": SYSTEM_FEW_SHOT_COT}]
    _append_few_shot_pairs(messages, shot_examples)
    messages.append(
        {
            "role": "user",
            "content": USER_FEW_SHOT_COT_TEMPLATE.format(question=question.question),
        }
    )
    return call_llm(messages)