#!/usr/bin/env python3
"""타깃 언어별 CJK/스크립트 오염 마스크를 만든다 — 언어마다 규칙이 다르다.

## 왜 언어마다 다른가 (실측이 강제했다)

한국어용 마스크(M2)를 일본어에 그대로 대 봤더니 상용어 12개 중 **11개가 잘렸다**
(日本語·会議·時間·電話·勉強·経済·国際·実際·学校·写真·広告). 한국어는 한자를 낱글자
글로스로만 쓰므로 "2자 이상 순수한자 토큰 = 중국어 단어"가 성립하지만, 일본어는 그런
복합어가 **모국어 정서법**이다. 규칙이 그대로 뒤집힌다.

## 무엇으로 "이 언어의 글자"를 판정하나

OpenCC 의 간체/번체 축은 일본어에서 **틀린다**. 신자체 国·学·会·写·独·当·来 를
`s2t` 가 전부 "간체"로 판정한다(우리 M1 이 일본어를 깨는 이유). 대신 이 스크립트는
**legacy 국가 인코딩을 그 나라의 문자 레퍼토리로** 쓴다.

    shift_jis   -> 일본 한자 레퍼토리 (신자체 포함, 중국 간체전용 제외)
    big5        -> 번체 레퍼토리
    big5hkscs   -> 번체 + 홍콩 확장 (광둥어 고유자 嘅喺啲哋嘢 포함)
    euc_kr      -> 한국 한자 레퍼토리
    iso8859-6   -> 표준 아랍어 자모 (페르소-우르두 확장자 제외)

⛔ **big5 로 광둥어를 처리하면 안 된다.** 광둥어 고유자 11자 중 5자(嘅喺啲哋嘢)가
plain Big5 밖이라, zh-TW 마스크를 그대로 쓰면 광둥어 자신이 잘린다. yue 는 big5hkscs.

## 티어

언어마다 티어 수가 다르다. 이것도 발견이지 게으름이 아니다 — "정당한 겹침"이 없는
언어(vi)는 T1=T2=T3 이고, 겹침이 전면적인 언어(ja)는 공격적 티어가 아예 못 쓸 물건이다.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent

HAN    = [(0x4E00, 0x9FFF), (0x3400, 0x4DBF), (0xF900, 0xFAFF)]
KANA   = [(0x3040, 0x309F), (0x30A0, 0x30FF), (0x31F0, 0x31FF)]
HANGUL = [(0xAC00, 0xD7A3), (0x1100, 0x11FF), (0x3130, 0x318F)]
ARABIC = [(0x0600, 0x06FF), (0x0750, 0x077F), (0x08A0, 0x08FF)]
ARPRES = [(0xFB50, 0xFDFF), (0xFE70, 0xFEFF)]   # 아랍어 표현형(레거시 리가처)

# 아랍어에서 iso8859-6 밖이지만 **정당한** 코드포인트 — 숫자·타슈킬·꾸란 기호
AR_KEEP = ([(0x0660, 0x066D)] +          # 아랍-인도 숫자 · 구두점
           [(0x064B, 0x065F)] +          # 타슈킬(모음부호)
           [(0x0670, 0x0671)] +          # superscript alef · alef wasla
           [(0x06D6, 0x06ED)])           # 꾸란 주석 기호


def inr(cp, rs):
    return any(a <= cp <= b for a, b in rs)


def enc_ok(c: str, e: str) -> bool:
    try:
        c.encode(e); return True
    except Exception:
        return False


# ── 언어별 규칙 ────────────────────────────────────────────────────────────
# cut(s) -> set of tier names this token belongs to.

def rules_ko(s, cps, s2t):
    """한국어: 낱글자 한자 병기는 살리고 중국어 단어 통짜를 자른다(기존 M2 와 동형)."""
    if any(inr(c, HANGUL) for c in cps):
        return set()
    han = [c for c in s if inr(ord(c), HAN)]
    kana = any(inr(c, KANA) for c in cps)
    if not han and not kana:
        return set()
    out = {"t3"}
    if kana or any(s2t.convert(c) != c for c in han):
        out |= {"t1", "t2"}
    elif len(han) >= 2 and len(s.strip()) == len(han):
        out |= {"t2"}
    return out


def rules_ja(s, cps, s2t):
    """일본어: 가나가 있으면 절대 보존. JIS 레퍼토리 밖 한자만 자른다."""
    if any(inr(c, KANA) for c in cps):
        return set()                                   # ⛔ 일본어 전용 스크립트
    han = [c for c in s if inr(ord(c), HAN)]
    out = set()
    if han and any(not enc_ok(c, "shift_jis") for c in han):
        out |= {"t1", "t2", "t3"}                      # 중국 간체전용 (这们说华长东…)
    if any(inr(c, HANGUL) for c in cps):
        out |= {"t2", "t3"}                            # 한글 누출
    if han and any(s2t.convert(c) != c for c in han):
        out |= {"t3"}                                  # ⛔ 순진한 이식 = 신자체 파괴
    return out


def _rules_trad(enc):
    def f(s, cps, s2t):
        han = [c for c in s if inr(ord(c), HAN)]
        out = set()
        if han and any(not enc_ok(c, enc) for c in han):
            out |= {"t1", "t2", "t3"}                  # 간체전용 + 신자체전용
        if any(inr(c, KANA) for c in cps):
            out |= {"t1", "t2", "t3"}
        if any(inr(c, HANGUL) for c in cps):
            out |= {"t2", "t3"}
        return out
    return f


def rules_vi(s, cps, s2t):
    """베트남어: 현대 정서법에 CJK 정당 사용이 0 — 겹침이 없으니 티어도 없다."""
    if any(inr(c, HAN) or inr(c, KANA) or inr(c, HANGUL) for c in cps):
        return {"t1", "t2", "t3"}
    return set()


def rules_ar(s, cps, s2t):
    """아랍어: 페르소-우르두 전용 자모를 자르고 타슈킬·아랍인도숫자는 보존한다."""
    out = set()
    for c in cps:
        if inr(c, ARABIC):
            ch = chr(c)
            if not enc_ok(ch, "iso8859-6") and not inr(c, AR_KEEP):
                out |= {"t1", "t2", "t3"}              # پ چ ژ گ ی ک ٹ ڈ ڑ ں ھ ے …
        elif inr(c, ARPRES):
            out |= {"t2", "t3"}                        # 레거시 표현형
        elif inr(c, HAN) or inr(c, KANA) or inr(c, HANGUL):
            out |= {"t3"}
    return out


LANGS = {
    "ko":    dict(fn=rules_ko,  tiers=["t1", "t2", "t3"], default="t2",
                  name="Korean", leak="Chinese / Japanese"),
    "ja":    dict(fn=rules_ja,  tiers=["t1", "t2", "t3"], default="t1",
                  name="Japanese", leak="Chinese (simplified-only)"),
    "zh-TW": dict(fn=_rules_trad("big5"),      tiers=["t1", "t2"], default="t1",
                  name="Traditional Chinese", leak="Simplified Chinese / Japanese shinjitai"),
    "yue":   dict(fn=_rules_trad("big5hkscs"), tiers=["t1", "t2"], default="t1",
                  name="Cantonese", leak="Simplified Chinese / Japanese shinjitai"),
    "vi":    dict(fn=rules_vi,  tiers=["t1"], default="t1",
                  name="Vietnamese", leak="Chinese / Japanese / Korean"),
    "ar":    dict(fn=rules_ar,  tiers=["t1", "t2", "t3"], default="t1",
                  name="Arabic", leak="Persian / Urdu orthography"),
}

# ── 프로브: 토크나이저 수준에서 코드가 판정한다 (GPU 불요) ──────────────────
# preserve = 이 언어의 정당한 표현. 하나라도 잘리면 그 마스크는 실패다.
# suppress = 오염원. 안 잘리면 마스크가 무력하다.
PROBES = {
 "ko":    dict(preserve=["개항(開港)", "債權", "새옹지마", "大韓民國", "目的"],
               suppress=["您的", "贵公司", "具体时间", "本次会议", "ありがとう"]),
 "ja":    dict(preserve=["日本語", "会議", "時間", "電話", "勉強", "経済", "国際",
                         "実際", "学校", "写真", "広告", "ありがとうございます"],
               suppress=["这个", "们", "说话", "中华", "东西", "问题", "谁", "买卖"]),
 "zh-TW": dict(preserve=["電腦", "資訊", "網路", "繁體", "臺灣", "軟體", "國學會"],
               suppress=["电脑", "资讯", "网络", "简体", "软件", "実際", "気持"]),
 "yue":   dict(preserve=["嘅", "係", "冇", "喺", "啲", "唔", "佢哋", "乜嘢", "電腦"],
               suppress=["电脑", "资讯", "网络", "简体", "実際"]),
 "vi":    dict(preserve=["Xin chào", "Việt Nam", "cảm ơn bạn", "công nghệ"],
               suppress=["中国", "日本語", "电脑", "ありがとう", "한국"]),
 "ar":    dict(preserve=["مرحبا", "كيف حالك", "اللغة العربية", "الْحَمْدُ", "١٢٣", "شكرا", "ﷺ", "ﷲ"],
               suppress=["پ", "چ", "ژ", "گ", "سلام خوبی", "کتاب", "ہے"]),
}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="Qwen/Qwen3.8-27B")
    ap.add_argument("--langs", nargs="+", default=list(LANGS))
    ap.add_argument("--outdir", default=str(HERE / "masks"))
    a = ap.parse_args()

    from transformers import AutoTokenizer
    from opencc import OpenCC
    s2t = OpenCC("s2t")
    tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True)
    n = max(tok.get_vocab().values()) + 1
    outdir = Path(a.outdir); outdir.mkdir(parents=True, exist_ok=True)

    decoded = []
    for tid in range(n):
        try:
            decoded.append(tok.decode([tid], skip_special_tokens=False))
        except Exception:
            decoded.append("")

    rc = 0
    for lang in a.langs:
        cfg = LANGS[lang]
        tiers = {t: [] for t in cfg["tiers"]}
        for tid, s in enumerate(decoded):
            if not s:
                continue
            hit = cfg["fn"](s, [ord(c) for c in s], s2t)
            for t in cfg["tiers"]:
                if t in hit:
                    tiers[t].append(tid)

        sets = {t: set(v) for t, v in tiers.items()}
        print(f"\n=== {lang} ({cfg['name']}) · 누출원: {cfg['leak']}")
        for t in cfg["tiers"]:
            star = " ← 기본" if t == cfg["default"] else ""
            print(f"  {t}: {len(tiers[t]):>7,} 토큰{star}")

        # ⛔ 프로브 — 코드가 판정한다
        pr = PROBES[lang]
        dflt = sets[cfg["default"]]
        def toks(w): return tok.encode(w, add_special_tokens=False)
        bad_keep = [w for w in pr["preserve"] if any(i in dflt for i in toks(w))]
        bad_cut  = [w for w in pr["suppress"] if not any(i in dflt for i in toks(w))]
        print(f"  보존 프로브 {len(pr['preserve'])-len(bad_keep)}/{len(pr['preserve'])}"
              f"  · 억제 프로브 {len(pr['suppress'])-len(bad_cut)}/{len(pr['suppress'])}")
        if bad_keep:
            known = lang == "ko"      # M2 가 문서화한 대가(목적·사고류 통짜 한자어)
            tag = "⚠️ 알려진 대가" if known else "⛔ 잘리면 안 되는데 잘림"
            print(f"  {tag}: {bad_keep}")
            if not known:
                rc = 1
        if bad_cut:
            print(f"  ⚠️ 잘려야 하는데 남음: {bad_cut}")

        for t in cfg["tiers"]:
            p = outdir / f"mask_{lang}_{t}.json"
            p.write_text(json.dumps({
                "name": f"{lang}-{t}", "target_language": lang, "model": a.model,
                "leak_source": cfg["leak"], "is_default": t == cfg["default"],
                "target_token_ids": sorted(tiers[t]), "n_target": len(tiers[t]),
            }, ensure_ascii=False), encoding="utf-8")
    return rc


if __name__ == "__main__":
    raise SystemExit(main())