# -*- 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 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 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) -> 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)
        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 reorder(sample: Sequence[Any], spec: SlotSpec) -> List[Any]:
    """把"上一份 cd"的值按**本次构建**的槽位序重排。

    原理（#31）：同一台机器上的多份样本，**取值集合相同**、只有顺序不同；
    所以只要样本与本次构建的"槽位数一致"，就能按"项 -> 槽"的对应关系做置换。
    """
    if len(sample) != spec.n_slots:
        raise ValueError(
            f"样本槽数 {len(sample)} != 本次构建槽数 {spec.n_slots}；"
            f"两者必须来自同一环境（见 PROGRESS.md #31）")
    out: List[Any] = []
    for (_i, _m, slot, n) in spec.assign:
        out.extend(sample[slot:slot + n])
    return out


# --------------------------------------------------------------------------- #
# 明文组装 + 加密                                                               #
# --------------------------------------------------------------------------- #
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("--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)
    spec = SlotSpec(seq, kinds, v0s)
    print(f"文件        : {args.path}")
    print(f"清单模块    : registry[{_mid}]   项数 {len(seq)}")
    print(f"槽位规格    : {spec.describe()}")
    print(f"XTEA 密钥   : {key}")

    if args.from_sample:
        sample = load_sample_cd(args.from_sample)
        try:
            cd = reorder(sample, spec)
        except ValueError as exc:
            print(f"❌ {exc}")
            return 1
        print(f"样本槽数    : {len(sample)}  -> 重排后 {len(cd)}")
    elif args.values:
        vals = json.load(open(args.values, encoding="utf-8"))
        cd = []
        for (_i, m, _slot, n) in spec.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())