# -*- coding: utf-8 -*-
"""提示词集中管理模块。

设计目的：把"模型该怎么判断"和"怎么调用模型"分开——本文件只关心提示词内容，
llm_client.py 只负责拿着提示词去调接口。以后要调整分类边界、增删细分类目，
只需要改这一个文件，不涉及调用逻辑。

两种模式共用同一套背景知识（什么是产品本体、资料/配件/耗材的定义），只是
模式为"fine"时，额外要求模型在第三行给出细分类目，且 few-shot 的回复也带
上第三行。

细分类目枚举以布莱迪真实历史数据（新参数字段，约2.4万条）为基础人工归纳，
并与业务方逐项确认边界后定稿：
- 螺栓/螺母单列"紧固件"一类（不并入法兰隔膜等附属部件）。
- 垫片单列第三大类"耗材/密封件"，不算在"配件"里，避免大量小垫片稀释"配件"
  标签的信息量（一张订单如果只要了一片垫片，和要了一整套阀组，不该被算作
  同等权重的"携带配件"）。
- 阻尼帽比照阻尼钉，视为工艺处理，不算配件、不进枚举。
- 合格证明确不计入判断：既不算资料，也不算任何其他附件类型。无论参数文本
  里是否提到合格证，一律忽略——不影响"是否携带附件"的判断结果，也不会
  出现在细分类目里。这是与此前版本的主要差异之一，few-shot 里同步给出了
  "仅有合格证"的反例（否则模型容易在缺少监督信号时把合格证当成资料的一种）。

"其他"兜底类目（可选，默认关闭）：
遇到确实是随货交付的附件、但不属于以上任何一个具体细分类目的情况，
如果调用方传入 enable_other=True，模型可以把这类附件归为第四大类"其他"，
细分类目也固定填"其他"（不允许模型自己编造具体名称，保持枚举严格校验的
风格不变）。enable_other=False（默认）时，"其他"完全不出现在提示词和
校验枚举里，行为与不支持"其他"之前完全一致——这个参数必须和
parser.parse_llm_output 使用的 enable_other 保持一致，否则会出现
"提示词允许但校验拒绝"或反过来的不一致；classifier.py 负责保证两边一致。
"""
from __future__ import annotations

# ============================================================
# 细分类目枚举（fine 模式专用，供 parser.py 做严格校验）
# ============================================================

DOC_FINE_TYPES: list[str] = [
    "出厂检测检验报告",
    "材质证明类",
    "说明书",
    "检验记录过程性",
    "图纸类",
    "标定校验证书",
    "质量证明书",
    "原产地证明",
    "检验合格证书综合",
    "其他交工文件",
    "营业执照",
    "型式检验报告",
]

PART_FINE_TYPES: list[str] = [
    "针型阀",
    "球阀",
    "截止阀",
    "旋塞阀",
    "角阀",
    "阀组",
    "冷凝圈",
    "冷凝管",
    "冷凝弯",
    "虹吸管",
    "表弯管",
    "缓冲管缓冲弯",
    "法兰隔膜",
    "隔离器",
    "过压保护器",
    "散热器散热片",
    "铅封",
    "转换接头",
    "焊接接头短节短管",
    "卡箍抱箍",
    "紧固件",
    "活接头",
    "接线盒",
    "电缆插头",
    "变送器",
]

CONSUMABLE_FINE_TYPES: list[str] = [
    "垫片",
]

# "其他"大类专用：整个大类只有这一个细分值，且值本身就是"其他"——不允许
# 模型编造具体名称，兜底之外仍保持和三大类同样严格的枚举校验风格。
OTHER_CATEGORY = "其他"
OTHER_FINE_TYPE = "其他"

# 三大类固定存在；"其他"是否作为第四类生效，完全由调用方传入的 enable_other
# 决定，默认不生效。
CATEGORY_NAMES = ("资料", "配件", "耗材")

# 三大类下所有细分类目的全集（不含"其他"）
ALL_FINE_TYPES: set[str] = set(DOC_FINE_TYPES) | set(PART_FINE_TYPES) | set(CONSUMABLE_FINE_TYPES)

# 细分类目 -> 所属大类（不含"其他"）
FINE_TYPE_TO_CATEGORY: dict[str, str] = {
    **{t: "资料" for t in DOC_FINE_TYPES},
    **{t: "配件" for t in PART_FINE_TYPES},
    **{t: "耗材" for t in CONSUMABLE_FINE_TYPES},
}


def valid_categories(enable_other: bool = False) -> set[str]:
    """返回当前生效的大类合法值集合，enable_other=True 时多一个"其他"。"""
    return set(CATEGORY_NAMES) | ({OTHER_CATEGORY} if enable_other else set())


def all_fine_types(enable_other: bool = False) -> set[str]:
    """返回当前生效的细分类目全集，enable_other=True 时多一个"其他"。"""
    return ALL_FINE_TYPES | ({OTHER_FINE_TYPE} if enable_other else set())


def fine_type_to_category(enable_other: bool = False) -> dict[str, str]:
    """返回当前生效的"细分类目 -> 大类"映射，enable_other=True 时多"其他"->"其他"。"""
    if not enable_other:
        return dict(FINE_TYPE_TO_CATEGORY)
    return {**FINE_TYPE_TO_CATEGORY, OTHER_FINE_TYPE: OTHER_CATEGORY}


# ============================================================
# 背景知识（两种模式共用）
# ============================================================

_BACKGROUND = """你是布莱迪压力表公司的订单参数审核助手。你的任务是判断一条订单的"技术参数"文本中，除了压力表产品本体之外，是否还要求交付"附件"。

【背景】公司生产压力表（用于化工行业）。参数文本里大部分内容是描述产品本体的规格，这些都不是附件。

【什么是产品本体（不算附件）】
表壳直径/材质、精度等级、防护等级(IP)、过压保护、弹性元件材质、接头材质、接头螺纹、安装方式、表盘材质/刻度/颜色、量程、充油、玻璃、指针、调零装置、缓冲钉/阻尼钉/阻尼帽（均属工艺处理，不算配件）、结构代号(如513/531)、卡盘尺寸、以及各种打标/示值/红线标识要求。

【附件 = 产品本体之外，随货交付的东西，分三大类】

一、资料类（交付给客户的技术文件）：
检测报告、检验报告、出厂报告、材质证明/材质单/材质报告、说明书(使用/安装/维护)、标定证书/校验证书、制造记录/检验记录、图纸、含技术信息的二维码/标签、营业执照、原产地证明、型式检验报告。
※ 注意：装箱单、送货单、贴箱单、发货章、包装/箱体标注等，属于物流包装指令，不算资料。合格证也不算资料——无论参数文本里是否提到合格证，一律忽略，既不影响"是否携带附件"的判断，也绝不能出现在第三行的具体分类里。

二、配件类（随货的独立实物零件，不含垫片）：
针型阀、球阀、截止阀、旋塞阀、角阀、阀组（二/三/五阀组等）、冷凝管/冷凝圈/冷凝弯、表弯管/缓冲弯、虹吸管、过压保护器、散热器/散热片、焊接接头/短节/短管、法兰隔膜、隔离器、铅封、转换接头、卡箍/抱箍、紧固件（螺栓/螺母）、活接头、接线盒、电缆/插头、变送器等，凡是随表一起出货的独立零件都算。
※ 注意：位号牌、铭牌、标牌、号牌属于标识件，不算配件。缓冲钉/阻尼钉/阻尼帽都是工艺处理，不算配件。

三、耗材/密封件类：
垫片（四氟垫片、紫铜垫片、缠绕垫、密封圈、O型圈等各种材质的垫片），单独列为耗材类，不算在配件里。

【判断步骤】
1. 逐项扫描文本，剔除所有"产品本体"描述，并剔除合格证（合格证不计入判断，见上方资料类说明）。
2. 看剩下的内容里是否出现资料类、配件类或耗材类的东西。
3. 只要出现任一类，第一行就是"是"；三类都没有就是"否"。"""

_OTHER_BACKGROUND_SUFFIX = """

四、其他类（兜底类目，仅在本次调用启用该选项时可用）：
如果确认订单确实随货交付了某个附件，但它明显不属于以上资料/配件/耗材已列出的任何一个具体细分类目（包括各大类里已有的"其他XX"兜底项，如"其他交工文件""检验合格证书综合"），也无法合理归入其中任何一类，才可以归为"其他"。"其他"是最后的兜底选项——凡是能归到前三类具体细分类目的，都必须归到那里，不要图省事就往"其他"里放。"""


def _background(enable_other: bool) -> str:
    return _BACKGROUND + (_OTHER_BACKGROUND_SUFFIX if enable_other else "")


# ============================================================
# coarse 模式：只输出大类
# ============================================================

def _coarse_system_prompt(enable_other: bool) -> str:
    ordered_categories = list(CATEGORY_NAMES) + ([OTHER_CATEGORY] if enable_other else [])
    categories_str = "、".join(f'"{c}"' for c in ordered_categories)
    return (
        _background(enable_other)
        + f"""

【输出格式 —— 必须严格按此格式，只输出这两行，不要多余文字、不要标点符号、不要解释】
是否携带附件: 是
附件类型: 资料,配件

【格式规则】
- 第一行只能是"是否携带附件: 是"或"是否携带附件: 否"。
- 第二行"附件类型: "后面，如果第一行是"否"，留空；如果是"是"，写{categories_str}中的一个或多个，用英文逗号分隔（不要空格）。
- 不要输出 JSON、不要输出 markdown 代码块、不要输出任何解释性文字。"""
    )


FEWSHOT_COARSE: list[tuple[str, str]] = [
    (
        "表壳直径:100mm 表壳材质:304SS 精度等级:1.6级 接头螺纹:M20×1.5 带缓冲钉 513结构",
        "是否携带附件: 否\n附件类型: ",
    ),
    (
        "表壳直径:63mm 精度等级:2.5级 带不锈钢位号牌",
        "是否携带附件: 否\n附件类型: ",
    ),
    (
        "表壳直径:100mm 张贴好装箱单并注明订单号 送货单盖发货章",
        "是否携带附件: 否\n附件类型: ",
    ),
    (
        "接头螺纹:ZG1/4 带阻尼帽 不要阻尼钉",
        "是否携带附件: 否\n附件类型: ",
    ),
    (
        "表壳直径:100mm 表壳材质:304SS 合格证随货",
        "是否携带附件: 否\n附件类型: ",
    ),
    (
        "表壳直径:100mm 提供出厂检测报告 合格证随货",
        "是否携带附件: 是\n附件类型: 资料",
    ),
    (
        "表壳直径:100mm 针型阀材质:304SS 配316L表弯管",
        "是否携带附件: 是\n附件类型: 配件",
    ),
    (
        "配四氟垫片 法兰专用四氟垫片",
        "是否携带附件: 是\n附件类型: 耗材",
    ),
    (
        "出厂检验报告 仪表使用说明书 配冷凝圈 过压保护器 配四氟垫片",
        "是否携带附件: 是\n附件类型: 资料,配件,耗材",
    ),
]

# 仅在 enable_other=True 时追加到 FEWSHOT_COARSE 后面
FEWSHOT_COARSE_OTHER: list[tuple[str, str]] = [
    (
        "表壳直径:100mm 随货附带一个定制不锈钢防雨罩",
        "是否携带附件: 是\n附件类型: 其他",
    ),
]


# ============================================================
# fine 模式：额外输出第三行细分类目
# ============================================================

def _fine_system_prompt(enable_other: bool) -> str:
    ordered_categories = list(CATEGORY_NAMES) + ([OTHER_CATEGORY] if enable_other else [])
    categories_str = "、".join(f'"{c}"' for c in ordered_categories)
    other_line = (
        f"\n其他类细分：{OTHER_FINE_TYPE}（仅当确认是随货附件、但资料/配件/耗材下所有具体"
        f"细分类目都无法覆盖时才使用，不需要也不允许编造更具体的名称）"
        if enable_other else ""
    )
    return (
        _background(enable_other)
        + f"""

【细分类目枚举 —— 第三行必须只从下面这些词里选，不能自己编造新词】
资料类细分：{"、".join(DOC_FINE_TYPES)}
配件类细分：{"、".join(PART_FINE_TYPES)}
耗材类细分：{"、".join(CONSUMABLE_FINE_TYPES)}{other_line}

【输出格式 —— 必须严格按此格式，只输出这三行，不要多余文字、不要标点符号、不要解释】
是否携带附件: 是
附件类型: 资料,配件
具体分类: 出厂检测检验报告,针型阀

【格式规则】
- 第一行只能是"是否携带附件: 是"或"是否携带附件: 否"。
- 第二行"附件类型: "后面，如果第一行是"否"，留空；如果是"是"，写{categories_str}中的一个或多个，用英文逗号分隔（不要空格）。
- 第三行"具体分类: "后面，如果第一行是"否"，留空；如果是"是"，从上面枚举的细分类目中选出所有识别到的，用英文逗号分隔（不要空格），不能选枚举之外的词，也不能只写大类名称（如"资料"）。
- 第三行选出的每个细分类目，都必须能对应到第二行里的某个大类（例如选了"针型阀"，第二行就必须包含"配件"）。
- 不要输出 JSON、不要输出 markdown 代码块、不要输出任何解释性文字。"""
    )


FEWSHOT_FINE: list[tuple[str, str]] = [
    (
        "表壳直径:100mm 表壳材质:304SS 精度等级:1.6级 接头螺纹:M20×1.5 带缓冲钉 513结构",
        "是否携带附件: 否\n附件类型: \n具体分类: ",
    ),
    (
        "表壳直径:63mm 精度等级:2.5级 带不锈钢位号牌",
        "是否携带附件: 否\n附件类型: \n具体分类: ",
    ),
    (
        "表壳直径:100mm 张贴好装箱单并注明订单号 送货单盖发货章",
        "是否携带附件: 否\n附件类型: \n具体分类: ",
    ),
    (
        "接头螺纹:ZG1/4 带阻尼帽 不要阻尼钉",
        "是否携带附件: 否\n附件类型: \n具体分类: ",
    ),
    (
        "表壳直径:100mm 表壳材质:304SS 合格证随货",
        "是否携带附件: 否\n附件类型: \n具体分类: ",
    ),
    (
        "表壳直径:100mm 提供出厂检测报告 合格证随货",
        "是否携带附件: 是\n附件类型: 资料\n具体分类: 出厂检测检验报告",
    ),
    (
        "表壳直径:100mm 针型阀材质:304SS 配316L表弯管",
        "是否携带附件: 是\n附件类型: 配件\n具体分类: 针型阀,表弯管",
    ),
    (
        "配四氟垫片 法兰专用四氟垫片",
        "是否携带附件: 是\n附件类型: 耗材\n具体分类: 垫片",
    ),
    (
        "带不锈钢螺栓螺母固定 配304SS法兰隔膜",
        "是否携带附件: 是\n附件类型: 配件\n具体分类: 紧固件,法兰隔膜",
    ),
    (
        "出厂检验报告 仪表使用说明书 配冷凝圈 过压保护器 配四氟垫片",
        "是否携带附件: 是\n附件类型: 资料,配件,耗材\n具体分类: 出厂检测检验报告,说明书,冷凝圈,过压保护器,垫片",
    ),
]

# 仅在 enable_other=True 时追加到 FEWSHOT_FINE 后面
FEWSHOT_FINE_OTHER: list[tuple[str, str]] = [
    (
        "表壳直径:100mm 随货附带一个定制不锈钢防雨罩",
        "是否携带附件: 是\n附件类型: 其他\n具体分类: 其他",
    ),
]


def get_prompt(mode: str, enable_other: bool = False) -> tuple[str, list[tuple[str, str]]]:
    """按模式（以及是否启用"其他"兜底类目）返回 (system_prompt, few_shot_examples)。

    mode 只接受 "coarse" 或 "fine"，其他值直接抛 ValueError——这里不做
    宽松兼容，避免拼写错误的 mode 静默退化成某个默认行为。

    enable_other 默认 False：不传该参数时，"其他"完全不出现在提示词里，
    模型和之后的 parser 校验都不知道这个选项的存在，行为与不支持"其他"
    之前完全一致。调用方（llm_client.py）必须把这里用的 enable_other 原样
    传给 parser.parse_llm_output，两边必须一致，否则会出现"提示词允许但
    校验拒绝"或反过来的不一致。
    """
    if mode == "coarse":
        prompt = _coarse_system_prompt(enable_other)
        fewshot = FEWSHOT_COARSE + (FEWSHOT_COARSE_OTHER if enable_other else [])
        return prompt, fewshot
    if mode == "fine":
        prompt = _fine_system_prompt(enable_other)
        fewshot = FEWSHOT_FINE + (FEWSHOT_FINE_OTHER if enable_other else [])
        return prompt, fewshot
    raise ValueError(f"不支持的分类模式: {mode!r}，只能是 'coarse' 或 'fine'")