"""
Few-shot CoT 教学演示入口。

运行前：
  pip install -r requirements.txt
  配置系统环境变量 DASHSCOPE_API_KEY

运行：
  python main.py              # 默认对比三种模式（各 1 次 API 调用）
  python main.py --shots 1    # Few-shot 只用 1 个示例（观察示例数量的影响）
  python main.py --shots 2    # Few-shot 用 2 个示例（默认）
"""

from __future__ import annotations

import argparse
import re
import sys

from config import API_KEY
from demo import run_direct, run_few_shot_cot, run_zero_shot_cot
from display import print_banner, print_intro, print_kv, print_model_reply, print_section
from examples import FEW_SHOT_EXAMPLES
from llm_client import get_call_count
from question import MAIN_QUESTION


def _configure_stdout() -> None:
    if hasattr(sys.stdout, "reconfigure"):
        try:
            sys.stdout.reconfigure(encoding="utf-8")
        except Exception:
            pass


def check_env() -> None:
    if not API_KEY:
        raise SystemExit(
            "未读取到系统环境变量 DASHSCOPE_API_KEY。\n"
            "请在 Windows「系统属性 → 环境变量」中配置后，重新打开终端再运行。"
        )


def extract_final_answer(text: str) -> str:
    match = re.search(r"最终答案[:：]\s*(.+)", text)
    if match:
        return match.group(1).strip()
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    return lines[-1] if lines else "（未识别到最终答案）"


def has_reasoning_steps(text: str) -> bool:
    markers = ("第一步", "第二步", "第三步", "年龄差", "最终答案")
    return sum(marker in text for marker in markers) >= 2


def is_answer_correct(final_answer: str, expected: str) -> bool:
    numbers = re.findall(r"\d+", expected)
    return all(num in final_answer for num in numbers)


def print_few_shot_preview(shot_count: int) -> None:
    print_section(f"本次 Few-shot CoT 将使用 {shot_count} 个示例")
    for example in FEW_SHOT_EXAMPLES[:shot_count]:
        print(f"\n  · {example.title}")
        print_kv("    教学点", example.teaching_point)


def run_demo(shot_count: int) -> None:
    question = MAIN_QUESTION
    few_shot_examples = FEW_SHOT_EXAMPLES[:shot_count]

    print_banner(question.title)
    print_kv("题目", question.question)
    print_kv("参考答案", question.expected_answer)
    print_kv("常见误区", question.common_mistake)
    print_kv("标准推理", question.correct_reasoning)
    print_few_shot_preview(shot_count)

    modes = [
        ("模式 A · 无 CoT（直接要答案）", lambda: run_direct(question)),
        (
            "模式 B · 零样本 CoT（只说「请一步一步思考」，不给示例）",
            lambda: run_zero_shot_cot(question),
        ),
        (
            f"模式 C · 少样本 CoT（先给 {shot_count} 个推理示例，再答本题）",
            lambda: run_few_shot_cot(question, few_shot_examples),
        ),
    ]

    results: list[tuple[str, str, bool, bool]] = []

    for step, (mode_name, runner) in enumerate(modes, start=1):
        print_section(mode_name)
        print(
            f"  [步骤 {step}/3] 正在调用百炼 API...\n"
            f"  流程：prompts.py / examples.py → demo.py 拼 messages → llm_client.py"
        )
        reply = runner()
        final_answer = extract_final_answer(reply)
        correct = is_answer_correct(final_answer, question.expected_answer)
        stepped = has_reasoning_steps(reply)
        print_model_reply(reply)
        print(f"\n  提取到的最终答案：{final_answer}")
        print(f"  答案是否正确：{'✓ 正确' if correct else '✗ 错误'}")
        print(f"  是否含分步推理：{'是' if stepped else '否'}")
        results.append((mode_name, final_answer, correct, stepped))

    print_section("三种模式对比小结")
    for mode_name, final_answer, correct, stepped in results:
        print(f"  · {mode_name}")
        print(f"    最终答案：{final_answer}（{'正确' if correct else '错误'}）")
        print(f"    分步推理：{'有' if stepped else '无'}")

    print(
        "\n  课堂观察要点：\n"
        "  · 模式 A 可能直接给出 105、175 等错误数字（把「一半」误用到现在）；\n"
        "  · 模式 B 常会展开推理，但步骤是否抓住「年龄差不变」不一定稳定；\n"
        "  · 模式 C 会模仿示例格式，更容易写出「先算年龄差 → 再算现在年龄」。\n"
        "  · 可尝试 --shots 1 与 --shots 2，观察示例数量对效果的影响。"
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Few-shot CoT 教学演示")
    parser.add_argument(
        "--shots",
        type=int,
        choices=range(1, len(FEW_SHOT_EXAMPLES) + 1),
        default=len(FEW_SHOT_EXAMPLES),
        metavar="N",
        help=f"Few-shot 示例数量（1-{len(FEW_SHOT_EXAMPLES)}），默认 {len(FEW_SHOT_EXAMPLES)}",
    )
    return parser.parse_args()


def main() -> None:
    _configure_stdout()
    check_env()
    args = parse_args()
    print_intro()
    run_demo(args.shots)

    print_banner("演示结束")
    print(f"本次共发起 {get_call_count()} 次百炼 API 调用。")
    print("题目位置：question.py（主题目）、examples.py（Few-shot 示例题）")
    print("完整说明：docs/项目说明.md")
    print("讲课话术：docs/课堂讲解指南.md")
    print("对比数学陷阱题：../cot_demo/")


if __name__ == "__main__":
    main()