#!/usr/bin/env python3
"""Few-shot prompt 构造器（无 key 可跑）。

按官方最佳实践（访问于 2026-08-30,
https://platform.claude.com/docs/en/build-with-claude/prompt-engineering/claude-prompting-best-practices）
构造 few-shot prompt:
  1. 示例相关 (Relevant)  2. 多样 (Diverse)  3. 用 <example>/<examples> 标签结构化
  4. 官方建议 3-5 个示例
并内置一个「多样性」粗检：标签分布不偏斜、无重复。

运行: python3 few_shot_builder.py
"""

from __future__ import annotations

from collections import Counter


class FewShotPromptBuilder:
    def __init__(self, task_instruction: str, role: str | None = None):
        self.role = role
        self.task_instruction = task_instruction
        self.examples: list[tuple[str, str]] = []

    def add_example(self, user_input: str, ideal_output: str) -> "FewShotPromptBuilder":
        self.examples.append((user_input, ideal_output))
        return self

    def diversity_report(self) -> dict:
        """粗粒度多样性检查：答案标签分布 + 重复输入检测。"""
        outputs = [out for _, out in self.examples]
        labels = Counter(
            next((w for w in ("refund", "bug", "feature", "howto") if w in out.lower()), "?")
            for out in outputs
        )
        inputs = [inp for inp, _ in self.examples]
        return {
            "n_examples": len(self.examples),
            "recommended_range": "3-5",
            "label_distribution": dict(labels),
            "duplicate_inputs": len(inputs) != len(set(inputs)),
        }

    def build(self) -> str:
        parts = []
        if self.role:
            parts.append(f"<role>{self.role}</role>")
        parts.append(f"<instructions>{self.task_instruction}</instructions>")
        ex_xml = []
        for i, (inp, out) in enumerate(self.examples, 1):
            ex_xml.append(
                f'<example index="{i}">\n  <input>{inp}</input>\n'
                f"  <ideal_output>{out}</ideal_output>\n</example>"
            )
        if ex_xml:
            parts.append("<examples>\n" + "\n".join(ex_xml) + "\n</examples>")
        # 长上下文排布原则：新输入（待处理数据）放最后
        parts.append("<input>{{NEW_USER_INPUT}}</input>")
        return "\n\n".join(parts)


if __name__ == "__main__":
    builder = (
        FewShotPromptBuilder(
            task_instruction="Classify each support email into exactly one category: "
            "refund | bug | feature | howto. Reply with the label only.",
            role="You are a senior support triage specialist.",
        )
        .add_example(
            "I was charged twice this month, please return my money.",
            "refund",
        )
        .add_example(
            "The export button crashes on Safari 17 every time.",
            "bug",
        )
        .add_example(
            "It would be great if you supported dark mode on reports.",
            "feature",
        )
        .add_example(
            "How do I invite a teammate to my workspace?",
            "howto",
        )
    )

    report = builder.diversity_report()
    print("[多样性检查]", report)
    assert report["n_examples"] == 4 and 3 <= report["n_examples"] <= 5
    assert not report["duplicate_inputs"]
    # 标签分布应覆盖全部类别且不过度偏斜
    assert len(report["label_distribution"]) == 4, "每个类别至少 1 个示例，避免分布偏斜"

    print("\n[生成的 prompt]\n")
    print(builder.build())
    print("\n要点: <examples> 包裹全部示例; 指令在前、待处理输入在末尾（官方排布建议）。")