# -*- coding: utf-8 -*-
"""
tdc_collectgen.py —— **按本次构建**生成 collect 明文（并加密）。

承 har/1ms_dianxuan/PROGRESS.md #29~#36:
  * #29/#31 槽位序**按构建随机化**，不能用缓存/旧模板；
  * #34/#35 槽位序 = 清单模块的 require 顺序，**可纯静态导出**；
  * #36  「清单项 → cd 槽」不是一一对应：`kind===2` 的项不进 cd，
          `kind===1` 的返回值是 JSON 碎片串、串内逗号会切出额外槽。

本模块做两件事:
  1. `reorder(template_cd, mapping, values)` —— 把"语义值"按**本次构建的槽位序**排成 cd；
  2. `build(...)` —— 组装成 `{"cd":[...],"sd":{...}}` 明文，并用该构建的 XTEA 密钥加密。

⚠️ 关于「语义值」从哪来（诚实说明，别把它当成已经解决的事）:
  目前有两种取值来源，都可用，但**完备性不同**：
    * `--values <json>`：由调用方提供「语义名 -> 值」字典（最稳，推荐）；
    * `--from-sample <cd json>`：用**另一份已解密的真机明文**做"同源重排"。
      这要求两份样本跑在**同一台机器/同一环境**（取值集合相同，只有顺序不同，见 #31），
      适用于"先拿一份真机明文当素材、再按本次构建重排"的验证场景。
  **仍需补齐**：纯算地把每个语义名对应的值算出来（多数是静态环境量，可配置；
  canvas/WebGL 指纹这类必须取自真实环境）。这一步未完成前，纯离线生成还不能保证过风控。

用法:
    # 用一份同机真机明文做重排, 按本次构建输出 collect
    .venv/bin/python tools/tdc_collectgen.py <tdc.js> \\
        --from-sample /tmp/real_cd.json --out /tmp/collect.b64
"""
from __future__ import annotations

import argparse
import base64
import json
import re
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))
if str(ROOT / "tools") not in sys.path:
    sys.path.insert(0, str(ROOT / "tools"))

from tcaptcha.crypto import xtea_encrypt_bytes  # noqa: E402
from tcaptcha.tdc_key import extract_xtea_key_pure  # noqa: E402
import tdc_layout  # noqa: E402
import tdc_program as tp  # noqa: E402


# --------------------------------------------------------------------------- #
# 槽位映射                                                                      #
# --------------------------------------------------------------------------- #
def value_sig(v: Any) -> str:
    """值的**结构签名**（用于在没有环境指纹时当次级锚点）。

    只描述"形状 + 常量"，抹掉随时间/会话变化的具体数值：
      1791736322 -> "int:epochish"     [[34,96,3284],...] -> "list:pts"
      "MacIntel" -> "str:MacIntel"     [360,360]          -> "list:2"
    易变槽（时间戳/轨迹/随机串）的签名在两份样本里**类型一致**，数值不同 —— 正好
    用来把它们彼此对应起来，再由调用方决定"沿用样本值"还是"重新生成"。
    """
    if v is None:
        return "null"
    if isinstance(v, bool):
        return f"bool:{v}"
    if isinstance(v, int):
        return "int:epochish" if 1_700_000_000 <= v <= 2_000_000_000 else f"int:{v}"
    if isinstance(v, float):
        return "float"
    if isinstance(v, str):
        if len(v) > 200:
            return "str:big"
        if _RANDSTR_RE.fullmatch(v):
            return "str:rand"
        return f"str:{v[:24]}"
    if isinstance(v, list):
        if v and all(isinstance(x, list) and len(x) == 3 for x in v):
            return "list:pts"
        return f"list:{len(v)}"
    return type(v).__name__


_RANDSTR_RE = re.compile(r"[A-Za-z0-9+/=]{20,}")


class _Unmapped:
    """哨兵：标记"这个目标槽没找到源槽"。

    为什么不能用 None：真机 cd 里本身就有 null 值（`out += "null"` 分支，
    见 PROGRESS.md #36），用 None 当占位符会把合法的 null 槽误判成"没对齐"。
    """

    __slots__ = ()

    def __repr__(self) -> str:      # pragma: no cover - 仅用于调试输出
        return "<unmapped>"


_UNMAPPED = _Unmapped()


def build_anchor_map(src: "SlotSpec", dst: "SlotSpec") -> List[Tuple[int, int]]:
    """把「源构建的槽」对齐到「目标构建的槽」，返回 [(src_slot, dst_slot), ...]。

    这是重排的核心。三层锚点，逐层降级（都必须**双向唯一**，否则宁可不对齐）:

      ① **环境语义指纹**（主锚点，跨构建稳定，实测多重集交集 23/24）
         源项 i 与目标项 j 的 env_fingerprint 相同 -> 对齐。
      ② **结构签名**（次锚点，用于常量槽 / 读不到环境的项）
         value_sig 相同 -> 对齐。
      ③ 剩余未对齐的 -> 由调用方显式处理（不猜）。

    为什么不能按项下标或模块 id 对齐（两个都已实测证伪）:
      * 项下标：两次构建的清单顺序**就是**不同的东西，按下标对齐 = 恒等映射；
      * 模块 id：编译器每次重新分配，两份构建交集 28 个模块里 id->语义 相同的
        只有 1/11（2026-10-12 实测）。
    """

    def bipartite(keys_src: List[Any], keys_dst: List[Any]) -> Dict[int, int]:
        """只在**双向唯一**时配对，避免多对一的歧义错位。"""
        from collections import defaultdict

        by_dst: Dict[Any, List[int]] = defaultdict(list)
        by_src: Dict[Any, List[int]] = defaultdict(list)
        for i, k in enumerate(keys_src):
            if k:
                by_src[k].append(i)
        for j, k in enumerate(keys_dst):
            if k:
                by_dst[k].append(j)
        pair: Dict[int, int] = {}
        for k, srcs in by_src.items():
            dsts = by_dst.get(k) or []
            if len(srcs) == 1 and len(dsts) == 1:
                pair[srcs[0]] = dsts[0]
        return pair

    pairs: Dict[int, int] = {}
    # 锚点优先级（2026-10-12 实测，4 份全新在线构建）：
    #   锚点                        跨构建唯一可配对
    #   ① 语义名(案例§5.3 指纹)      35/37   <- 主锚点
    #   ② 静态语义指纹 fp           29/37
    #   ③ 指令直方图 hist           20/37
    #   ④ 结构签名 value_sig        （仅双方都有实测值时）
    # 所以：语义名先用，再逐级降级。每一级都要求**双向唯一**才配对。

    # ① 语义名（案例 §5.3：字符串指纹 + 结构兜底 + 重名消歧）
    pairs.update(bipartite([n or "" for n in src.names],
                           [n or "" for n in dst.names]))
    # ②③ 降级：静态语义指纹 -> 指令直方图
    for key in (lambda s: [f or "" for f in s.fps],
                lambda s: [h or "" for h in s.hists]):
        for i, j in bipartite(key(src), key(dst)).items():
            pairs.setdefault(i, j)
    # ④ 末级：结构签名（仅在双方都有实测值时可用）
    sig_pair = bipartite([value_sig(v) for v in src.sample_values],
                         [value_sig(v) for v in dst.sample_values])
    for i, j in sig_pair.items():
        pairs.setdefault(i, j)

    out: List[Tuple[int, int]] = []
    for i, (_i, _m, sslot, n) in enumerate(src.assign):
        j = pairs.get(i)
        if j is None or j >= len(dst.assign):
            continue
        _dj, _dm, dslot, dn = dst.assign[j]
        if dn != n:
            continue                      # 槽数形状不一致 -> 拒绝猜，跳过
        for k in range(n):
            out.append((sslot + k, dslot + k))
    return out


def reorder(sample: Sequence[Any], dst: "SlotSpec",
            src: Optional["SlotSpec"] = None,
            fill: Optional[Dict[int, Any]] = None) -> List[Any]:
    """把「源样本」的 cd 值重排成**目标构建**的槽位序。

    ⚠️ 旧实现的 bug（2026-10-12 修复）：
        旧版签名是 `reorder(sample, spec)`，其中 `spec` 只是**目标构建**的规格，
        而循环里 `sample[slot:slot+n]` 用的 `slot` 恰恰是**目标构建自己的槽下标**。
        两次构建的槽位序不同（#31），所以这等价于"按目标下标切源数组" —— 顺序不同
        时全盘错位，顺序相同时才碰巧对。**必须同时给出源构建的规格**才能对齐。

    参数:
      sample - 源构建产出的 cd 值数组（真机明文）
      dst    - 目标构建的槽位规格（本次要发包的那个构建）
      src    - 源构建的槽位规格。**不传则退化为旧行为**（仅当两份构建同序时才正确，
              此时本函数会先校验再放行，不再静默错位）。
      fill   - {目标槽号: 值}，用于补齐锚点匹配不上的槽。

    锚点能力上限（2026-12 实测，见 tdc_layout.static_fingerprint / opcode_histogram）：
      静态锚点组合 `fp+hist` 单构建内唯一 **29/37**，跨构建交集 35。
      剩下 8 项（分成 2 组）**静态不可区分**：它们的函数体只有一处 JUMP 目标不同，
      语义藏在跳转目标里，符号求值取不到。这些槽必须由调用方用 `fill` 显式给值。

    未给 fill 且仍对不齐的槽 -> 抛 `UnmappedSlotsError`（宁可暴露"不知道"，
    也不要静默填错值 —— 旧实现正是静默填错，导致多轮实验建立在坏数据上）。
    """
    if src is None:
        # 契约：不提供源构建规格 = 调用方**承诺**样本与目标构建同序。
        # 此时重排只能是恒等映射；先把槽数校验掉，再直接返回副本。
        if len(sample) != dst.n_slots:
            raise ValueError(
                f"样本槽数 {len(sample)} != 目标构建槽数 {dst.n_slots}；"
                f"未提供源构建规格时不跨构建重排（见 PROGRESS.md #31）")
        out: List[Any] = [_UNMAPPED] * dst.n_slots
        for d in range(dst.n_slots):
            out[d] = sample[d]
        for d, v in (fill or {}).items():
            if 0 <= d < dst.n_slots:
                out[d] = v
        return out

    if len(sample) != src.n_slots:
        raise ValueError(
            f"样本槽数 {len(sample)} != 源构建槽数 {src.n_slots}；"
            f"样本必须来自 `src` 描述的那个构建")

    pairs = build_anchor_map(src, dst)
    # 用哨兵而不是 None 标记"未映射"：真机 cd 里本身就有 null 值，
    # 拿 None 当占位符会把合法的 null 槽误判成"没对齐"。
    out: List[Any] = [_UNMAPPED] * dst.n_slots
    for s_slot, d_slot in pairs:
        out[d_slot] = sample[s_slot]
    unmapped = [d for d, v in enumerate(out) if v is _UNMAPPED]
    for d, v in (fill or {}).items():
        if 0 <= d < dst.n_slots:
            out[d] = v
            if d in unmapped:
                unmapped.remove(d)
    if unmapped:
        raise UnmappedSlotsError(unmapped, dst, src)
    return out


class UnmappedSlotsError(RuntimeError):
    """有目标槽找不到源槽与之对应（锚点匹配失败）。"""

    def __init__(self, slots: List[int], dst: "SlotSpec", src: "SlotSpec") -> None:
        self.slots = slots
        super().__init__(
            f"目标构建有 {len(slots)} 个槽无法与源样本对齐（槽号 {slots[:12]}"
            f"{'...' if len(slots) > 12 else ''}）。"
            f"锚点不足时宁可报错也不要静默填错值 —— 请用 `--values` 显式提供这些槽的取值，"
            f"或检查两份构建是否来自同一采集器集合。")


def outer_commas(s: str) -> int:
    """数**不在**任何 `[]` / `{}` 内的逗号个数（方括号内的逗号不构成顶层分隔）。"""
    depth = 0
    n = 0
    for ch in s:
        if ch in "[{":
            depth += 1
        elif ch in "]}":
            depth = max(0, depth - 1)
        elif ch == "," and depth == 0:
            n += 1
    return n


class SlotSpec:
    """一次构建下的槽位规格（规则见 PROGRESS.md #36）。"""

    def __init__(self, seq: Sequence[int],
                 kinds: Optional[Sequence[Any]] = None,
                 v0s: Optional[Sequence[str]] = None,
                 fps: Optional[Sequence[Optional[str]]] = None,
                 hists: Optional[Sequence[str]] = None,
                 names: Optional[Sequence[Optional[str]]] = None) -> None:
        self.seq = list(seq)
        self.kinds = list(kinds) if kinds is not None else [None] * len(seq)
        # v0s[i] = 该项 get()[0] 的字符串化结果；kind=1 时按它的逗号数切出额外槽
        self.v0s = list(v0s) if v0s is not None else [""] * len(seq)
        # fps[i] = 该项的**环境语义指纹**（跨构建稳定, 见 tdc_layout.env_fingerprint）；
        # None/空串 表示该项读不到环境路径（常量槽 / 轨迹槽），需退到值签名或直接填值
        self.fps: List[Optional[str]] = list(fps) if fps is not None else [None] * len(seq)
        # hists[i] = 该项的**指令直方图**（兜底锚点，覆盖 37/37，见 tdc_layout.opcode_histogram）
        self.hists: List[str] = list(hists) if hists is not None else [""] * len(seq)
        # names[i] = 该项的**语义名**——案例 §5.3 的字符串+结构指纹（主锚点）。
        #   实测 4 份全新在线构建：37/37 识别、0 unknown、跨构建完全一致、**35/37 唯一可配对**。
        #   明显优于本项目早期自创的 fp/hist 锚点（29/37），所以放在最前面用。
        self.names: List[Optional[str]] = (
            list(names) if names is not None else [None] * len(seq))
        # 该构建下的一份**实测 cd 值**（源侧重排时要用它算结构签名）。
        # 目标侧可以留空；只有当成"源规格"参与 build_anchor_map 时才必须给。
        self.sample_values: List[Any] = []
        self.assign: List[Tuple[int, int, int, int]] = []   # (item, module, first_slot, n_slots)
        self._build()

    def _build(self) -> None:
        slot = 0
        for i, m in enumerate(self.seq):
            k = self.kinds[i]
            n = 1
            if k == 1 and isinstance(self.v0s[i], str):
                # 原样拼 -> 该串会被 JSON 解析器按逗号切成多个值。
                # 但**数组/对象形态**的串（'[]' / '[360,360]' / '["zh-CN","zh"]'）里的逗号
                # 处在方括号内, 拼回去仍是一个值, 不产生新槽 —— 只数"括号外"的逗号。
                n += outer_commas(self.v0s[i])
            if k == 2:
                # kind=2 把 r3 清空转存 r4（不接在 cd 后面），但 r4 会在收尾 @4905
                # 被二次加工后重新拼进 r3 —— 实测它**仍然产出一个槽**。
                n = 1
            self.assign.append((i, m, slot, n))
            slot += n

    @property
    def n_slots(self) -> int:
        return sum(a[3] for a in self.assign)

    def describe(self) -> str:
        multi = [(i, m, n) for (i, m, _s, n) in self.assign if n > 1]
        return (f"项 {len(self.seq)} -> 槽 {self.n_slots}  "
                f"(kind=2 跳过 {sum(1 for k in self.kinds if k == 2)} 项"
                + (f"; 多槽项 {multi}" if multi else "") + ")")


def load_shape(items_path: str, seq: Sequence[int]) -> Tuple[List[Any], List[str]]:
    """从运行时取样(items json)里读出每项的 kind 与 get()[0] 的字符串化值。"""
    d = json.load(open(items_path, encoding="utf-8"))
    res = d.get("result", d)
    rows = res.get("out") or res.get("items") or []
    kinds: List[Any] = []
    v0s: List[str] = []
    for r in rows:
        if not isinstance(r, dict):
            kinds.append(None); v0s.append(""); continue
        kinds.append(r.get("kind"))
        v = r.get("v0")
        v0s.append(str(v) if v is not None else "")
    if len(kinds) < len(seq):
        pad = len(seq) - len(kinds)
        kinds += [None] * pad
        v0s += [""] * pad
    return kinds[: len(seq)], v0s[: len(seq)]


def load_sample_cd(path: str) -> List[Any]:
    raw = json.load(open(path, encoding="utf-8"))
    return raw["cd"] if isinstance(raw, dict) else raw


def spec_from_js(path: str, kinds_path: Optional[str] = None,
                 sample_cd: Optional[Sequence[Any]] = None) -> "SlotSpec":
    """从一份 tdc.js 直接构造该构建的槽位规格（含环境语义指纹）。"""
    prog = tp.load(path)
    _mid, _entry, seq = tdc_layout.find_manifest_module(prog)
    kinds, v0s = load_shape(kinds_path, seq) if kinds_path else (None, None)
    fps = tdc_layout.item_fingerprints(prog, seq)
    hists = tdc_layout.item_histograms(prog, seq)
    # 主锚点：案例 §5.3 的字符串+结构指纹语义名（实测 35/37 唯一可配对）
    # 复用已加载的 prog（analyze_prog），避免把同一份 tdc.js 再解析一遍。
    # 不做静默降级（AGENTS.md 附录 A#8）：锚点算不出来就抛错，
    # 让调用方明确知道"这次没有主锚点"，而不是悄悄退化到较弱的锚点上。
    import tdc_modules
    names = [r["name"] for r in tdc_modules.analyze_prog(prog)]
    spec = SlotSpec(seq, kinds, v0s, fps, hists, names)
    if sample_cd is not None:
        spec.sample_values = list(sample_cd)
    return spec


def spec_from_seq(seq: Sequence[int],
                  kinds: Optional[Sequence[Any]] = None,
                  names: Optional[Sequence[Optional[str]]] = None,
                  sample_cd: Optional[Sequence[Any]] = None,
                  v0s: Optional[Sequence[str]] = None) -> "SlotSpec":
    """只用清单序列构造规格（**不需要 tdc.js**，供测试/合成用例使用）。

    纯算重排只需要「每一项占几槽」+「每项的锚点名」；
    tdc.js 只用来**导出**这两样东西，本身不是重排的输入。
    """
    spec = SlotSpec(seq, kinds, v0s, names=names)
    if sample_cd is not None:
        spec.sample_values = list(sample_cd)
    return spec


# --------------------------------------------------------------------------- #
# 明文组装 + 加密                                                               #
# --------------------------------------------------------------------------- #
def build_plain(cd: Sequence[Any], sd: Optional[dict] = None) -> str:
    sd = sd or {"od": "C", "ft": "qf_7Pf__H"}
    return json.dumps({"cd": list(cd), "sd": sd}, separators=(",", ":"))


def encrypt_plain(plain: str, key: Sequence[int]) -> str:
    data = plain.encode("latin1")
    data += b"\x00" * ((8 - len(data) % 8) % 8)      # 与真实样本一致: \x00 补齐
    return base64.b64encode(xtea_encrypt_bytes(data, list(key))).decode("ascii")


def main() -> int:
    ap = argparse.ArgumentParser(description="按本次构建生成 collect")
    ap.add_argument("path", help="本次会话的 tdc.js（**目标构建**）")
    ap.add_argument("--src-build", default=None,
                    help="源样本所属的那份 tdc.js。**跨构建重排必须给**；"
                         "不给则假定源样本与目标构建同序（并会校验槽数）")
    ap.add_argument("--from-sample", default=None, help="一份已解密的 cd 数组(json)")
    ap.add_argument("--values", default=None, help="语义名->值的 json 字典(实验性)")
    ap.add_argument("--out", default=None, help="输出 collect(base64) 的文件")
    ap.add_argument("--plain-out", default=None, help="同时输出明文")
    ap.add_argument("--kinds", default=None,
                    help="运行时取样的 items json（提供每项的 kind，见 #36）")
    args = ap.parse_args()

    prog = tp.load(args.path)
    _mid, _entry, seq = tdc_layout.find_manifest_module(prog)
    key = extract_xtea_key_pure(open(args.path, encoding="utf-8", errors="replace").read())["key"]

    kinds, v0s = load_shape(args.kinds, seq) if args.kinds else (None, None)
    fps = tdc_layout.item_fingerprints(prog, seq)
    hists = tdc_layout.item_histograms(prog, seq)
    names = [r["name"] for r in tdc_modules.analyze_prog(prog)]
    dst = SlotSpec(seq, kinds, v0s, fps, hists, names)
    anchored = sum(1 for f in fps if f)
    print(f"文件        : {args.path}")
    print(f"清单模块    : registry[{_mid}]   项数 {len(seq)}")
    print(f"槽位规格    : {dst.describe()}")
    print(f"锚点覆盖    : 语义名 {sum(1 for n in names if n and n != 'unknown')}/{len(seq)}"
          f"   静态指纹 {anchored}/{len(seq)}")
    print(f"XTEA 密钥   : {key}")

    if args.from_sample:
        sample = load_sample_cd(args.from_sample)
        if args.src_build:
            src = spec_from_js(args.src_build, args.kinds, sample)
        else:
            src = None
        try:
            cd = reorder(sample, dst, src)
        except (ValueError, UnmappedSlotsError) as exc:
            print(f"❌ {exc}")
            return 1
        print(f"样本槽数    : {len(sample)}  -> 重排后 {len(cd)}")
        if args.src_build:
            # 自证：锚点对齐后，值集合应当守恒（易变槽数值允许不同）
            print(f"对齐方式    : 环境语义指纹 + 结构签名（源={args.src_build}）")
    elif args.values:
        vals = json.load(open(args.values, encoding="utf-8"))
        cd = []
        for (_i, m, _slot, n) in dst.assign:
            v = vals[str(m)]
            cd.extend(v[:n] if isinstance(v, list) else [v])
        print(f"按语义字典  : {len(cd)} 槽")
    else:
        print("需要 --from-sample 或 --values")
        return 1

    plain = build_plain(cd)
    b64 = encrypt_plain(plain, key)
    print(f"明文长度    : {len(plain)}")
    print(f"collect     : {len(b64)} 字符 (tlg={len(b64)})")
    if args.plain_out:
        Path(args.plain_out).write_text(plain, encoding="latin1")
        print(f"明文 -> {args.plain_out}")
    if args.out:
        Path(args.out).write_text(b64, encoding="ascii")
        print(f"collect -> {args.out}")
    return 0


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