"""
CoT 三种模式的核心演示逻辑。

本文件负责「组装 messages + 调用模型」：
  build_*_messages()  →  拼提示词（可打印对比）
  run_*()             →  发请求并返回 messages + 回复
"""

from __future__ import annotations

from dataclasses import dataclass

from llm_client import call_llm
from prompts import (
    FEW_SHOT_EXAMPLE_ASSISTANT,
    FEW_SHOT_EXAMPLE_USER,
    SYSTEM_DIRECT,
    SYSTEM_FEW_SHOT_COT,
    SYSTEM_ZERO_SHOT_COT,
    USER_DIRECT_TEMPLATE,
    USER_FEW_SHOT_COT_TEMPLATE,
    USER_ZERO_SHOT_COT_TEMPLATE,
)
from questions import DemoQuestion

@dataclass
class ModeResult:
    """单次模式运行的完整结果，含提示词与模型回复。"""

    messages: list[dict[str, str]]
    reply: str
    prompt_summary: str


def build_direct_messages(question: DemoQuestion) -> list[dict[str, str]]:
    """无 CoT：system + user，只要求直接给答案。"""
    return [
        {"role": "system", "content": SYSTEM_DIRECT},
        {"role": "user", "content": USER_DIRECT_TEMPLATE.format(question=question.question)},
    ]


def build_zero_shot_cot_messages(question: DemoQuestion) -> list[dict[str, str]]:
    """零样本 CoT：system + user，要求分步思考，不给示例。"""
    return [
        {"role": "system", "content": SYSTEM_ZERO_SHOT_COT},
        {
            "role": "user",
            "content": USER_ZERO_SHOT_COT_TEMPLATE.format(question=question.question),
        },
    ]


def build_few_shot_cot_messages(question: DemoQuestion) -> list[dict[str, str]]:
    """少样本 CoT：system + 示例问答 + 新题目。"""
    return [
        {"role": "system", "content": SYSTEM_FEW_SHOT_COT},
        {"role": "user", "content": FEW_SHOT_EXAMPLE_USER},
        {"role": "assistant", "content": FEW_SHOT_EXAMPLE_ASSISTANT},
        {
            "role": "user",
            "content": USER_FEW_SHOT_COT_TEMPLATE.format(question=question.question),
        },
    ]


def run_direct(question: DemoQuestion) -> ModeResult:
    messages = build_direct_messages(question)
    reply = call_llm(messages)
    return ModeResult(
        messages=messages,
        reply=reply,
        prompt_summary="仅 system + user；要求只输出一行最终答案，禁止推理过程",
    )


def run_zero_shot_cot(question: DemoQuestion) -> ModeResult:
    messages = build_zero_shot_cot_messages(question)
    reply = call_llm(messages)
    return ModeResult(
        messages=messages,
        reply=reply,
        prompt_summary="system + user；user 中追加「请一步一步思考」，无示例",
    )


def run_few_shot_cot(question: DemoQuestion) -> ModeResult:
    messages = build_few_shot_cot_messages(question)
    reply = call_llm(messages)
    return ModeResult(
        messages=messages,
        reply=reply,
        prompt_summary="system + 示例user + 示例assistant + 新user；模型模仿示例推理格式",
    )