#!/usr/bin/env python3
"""合併標註結果、檢查標籤白名單，並和規則判定比對。

用法：merge_labels.py <dialogues.jsonl> <標註輸出目錄> <合併後 jsonl> [--review-dir DIR]

- 標註的 speaker 只有 confidence=high 才採用；unknown、none（非說出口的話）與其他信心都視為未判定。
- 標籤不在 tags.json 白名單內的整筆丟棄並列出。
- 以規則判定為高信心的句子，和標註比對一致率，用來估計兩邊的可靠度。
- --review-dir：每個角色輸出 20 句抽查清單（markdown），供人工檢查歸屬。
"""
import argparse
import json
import random
from collections import Counter
from pathlib import Path

TAGS = json.loads(Path(__file__).with_name("tags.json").read_text(encoding="utf-8"))
KEYS = {"id", "speaker", "confidence", "function", "emotion", "evidence"}


def read_jsonl(path):
    return [json.loads(line) for line in Path(path).read_text(encoding="utf-8").splitlines() if line.strip()]


# 有的批次輸出把 confidence 縮寫成一個字母；明確對應才轉換，其他值照常丟棄。
CONFIDENCE_ALIASES = {"h": "high", "m": "medium", "l": "low"}


def normalize(label):
    if label.get("confidence") in CONFIDENCE_ALIASES:
        return {**label, "confidence": CONFIDENCE_ALIASES[label["confidence"]]}
    return label


def check(label):
    errs = []
    if set(label) != KEYS:
        errs.append(f"欄位不符：{sorted(set(label) ^ KEYS)}")
    if label.get("function") not in TAGS["functions"]:
        errs.append(f"function 不在白名單：{label.get('function')}")
    if label.get("emotion") not in TAGS["emotions"]:
        errs.append(f"emotion 不在白名單：{label.get('emotion')}")
    if label.get("confidence") not in ("high", "medium", "low"):
        errs.append(f"confidence 無效：{label.get('confidence')}")
    return errs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("dialogues")
    ap.add_argument("labels_dir", help="標註輸出目錄；可用逗號接多個，後面的覆蓋前面的（用於重標）")
    ap.add_argument("out")
    ap.add_argument("--review-dir")
    a = ap.parse_args()
    recs = {r["id"]: r for r in read_jsonl(a.dialogues)}
    labels, rejected = {}, []
    files = [f for d in a.labels_dir.split(",") for f in sorted(Path(d).glob("*.jsonl"))]
    for f in files:
        for lb in read_jsonl(f):
            lb = normalize(lb)
            errs = check(lb)
            if lb.get("id") not in recs:
                errs.append("id 不存在")
            if errs:
                rejected.append((lb.get("id"), errs))
            else:
                labels[lb["id"]] = lb
    merged = []
    for rid, lb in labels.items():
        r = recs[rid]
        final = lb["speaker"] if lb["confidence"] == "high" and lb["speaker"] not in ("unknown", "none") else None
        merged.append({**r, "label": lb, "final_speaker": final})
    Path(a.out).write_text("".join(json.dumps(m, ensure_ascii=False) + "\n" for m in merged), encoding="utf-8")

    agree = Counter()
    for m in merged:
        r = m
        if r["confidence"] in ("high", "medium"):
            agree[(r["confidence"], r["speaker"] == m["label"]["speaker"])] += 1
    print(f"標註 {len(labels)} 句，丟棄 {len(rejected)} 筆")
    for rid, errs in rejected[:10]:
        print("  丟棄", rid, errs)
    print("標註採用（high 且非 unknown）：", sum(1 for m in merged if m["final_speaker"]), "/", len(merged))
    print("標註信心：", dict(Counter(m["label"]["confidence"] for m in merged)))
    print("非說出口的「」（none）：", sum(1 for m in merged if m["label"]["speaker"] == "none"))
    for conf in ("high", "medium"):
        ok, bad = agree[(conf, True)], agree[(conf, False)]
        if ok + bad:
            print(f"規則 {conf} 與標註一致：{ok}/{ok + bad}")

    if a.review_dir:
        out = Path(a.review_dir)
        out.mkdir(parents=True, exist_ok=True)
        rng = random.Random(11)
        for who in ("asuna", "suguha"):
            pool = [m for m in merged if m["final_speaker"] == who]
            rows = rng.sample(pool, min(20, len(pool)))
            lines = [f"# 抽查：{who}（{len(pool)} 句採用，抽 {len(rows)}）", "", "每句請看：說話者是否正確？不對的請在行尾標記。", ""]
            for m in rows:
                ctx = (m["before"] or m["prev_par"])[-60:]
                lines += [f"- `{m['id']}`「{m['text'][:60]}」", f"  - 前文：{ctx}", f"  - 標註：{m['label']['function']}／{m['label']['emotion']}；依據：{m['label']['evidence']}"]
            (out / f"{who}.md").write_text("\n".join(lines) + "\n", encoding="utf-8")
        print("抽查清單 ->", out)


if __name__ == "__main__":
    main()