# -*- coding: utf-8 -*-
"""
tdc_slotmap.py —— 把「清单项」精确映射到「cd 槽号」。

为什么需要单独一层（承 har/1ms_dianxuan/PROGRESS.md #34/#35/#36）:
  #34 槽位序 = 清单模块的 require 顺序（追加式数组，已核）
  #36 **槽数 ≠ 项数**: 装配循环里 `kind=1` 的分支是 `out += v`（不加引号直接拼），
      所以某些项的返回值**本身带逗号**时会一次性产出多个槽。
      实测同一构建: 37 项 -> 38 槽，多出的那一处是第 33 项（模块 1）
      get() 返回 "641547489,1791732001"。

因此映射不能是"第 i 项 = 第 i 槽"，必须按每个采集器的**返回值形状**推算:

    kind = ret[2]（可能缺省）
    kind == 1:  ret[0] 原样拼 -> 占 1 + ret[0].count(",") 槽（数字则 1 槽）
    其它:       ret[0] 占 1 槽
    ret[1] 是 number: 再 +1 槽（实测 kind=2 的项第二元素是 null，不触发）

本模块提供:
  * `slot_map(prog, items)` —— 给静态导出的项 + 运行时取样到的返回值形状，算出映射
  * `from_runtime_items(prog)` —— 需要浏览器环境（iframe 里读清单各项 get()），见 CLI
  * CLI 校验模式: 与一份已解密的 collect 明文对拍，确认推算出的槽数 == 实际槽数

用法:
    # 用运行时取样 + 本次 getData 解密出的明文, 做一次端到端对拍
    .venv/bin/python tools/tdc_slotmap.py <tdc.js> \\
        --items /tmp/items.json --plain /tmp/plain.txt
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import 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"))

import tdc_layout  # noqa: E402
import tdc_program as tp  # noqa: E402


def slots_of_item(ret) -> int:
    """一个采集器的 get() 返回值会占几个 cd 槽。

    ret 是运行时取到的值（或其形状描述）。规则见 PROGRESS.md #36。
    """
    if ret is None:
        return 1
    if not isinstance(ret, (list, tuple)) or not ret:
        return 1
    first = ret[0]
    kind = ret[2] if len(ret) >= 3 else None
    n = 1
    if kind == 1 and isinstance(first, str):
        # 不加引号直接拼 -> 串里的逗号会切出额外槽
        n += first.count(",")
    # 第二个元素是 number 时再占一槽（实测 kind=2 的第二元素是 null，不触发）
    if len(ret) >= 2 and isinstance(ret[1], (int, float)) and not isinstance(ret[1], bool):
        n += 1
    return n


def slot_map(seq: Sequence[int], rets: Sequence) -> List[Tuple[int, int, int]]:
    """返回 [(清单项下标, 模块 id, 该首槽在 cd 中的槽号)]。

    某一项占多槽时，只给它的**首槽**；其余槽是你按 ret 内容自然拼出来的。
    """
    out: List[Tuple[int, int, int]] = []
    slot = 0
    for i, m in enumerate(seq):
        out.append((i, m, slot))
        slot += slots_of_item(rets[i] if i < len(rets) else None)
    return out


def total_slots(seq: Sequence[int], rets: Sequence) -> int:
    return sum(slots_of_item(rets[i] if i < len(rets) else None) for i in range(len(seq)))


def load_items(path: str) -> List:
    """读 cdp_eval 的取样输出（{items:[{head,types,len,...}]}）。

    运行时只能带回 head（截断过），所以这里按 types/head 重建一个**够用的**形状。
    """
    d = json.load(open(path, encoding="utf-8"))
    res = d.get("result", d)
    items = res.get("items") or []
    rets = []
    for it in items:
        types = it.get("types")
        head = it.get("head")
        if not isinstance(types, list):
            rets.append(None)
            continue
        rebuilt = []
        for t, h in zip(types, head or []):
            if t == "string":
                rebuilt.append(h if isinstance(h, str) else str(h))
            elif t == "number":
                rebuilt.append(h if isinstance(h, (int, float)) else 0)
            elif t == "null":
                rebuilt.append(None)
            elif t == "array":
                rebuilt.append([])
            else:
                rebuilt.append(None)
        rets.append(rebuilt)
    return rets


def main() -> int:
    ap = argparse.ArgumentParser(description="清单项 -> cd 槽号的精确映射")
    ap.add_argument("path")
    ap.add_argument("--items", default=None, help="运行时取样的 items.json")
    ap.add_argument("--plain", default=None, help="已解密的 collect 明文（用于对拍校验）")
    ap.add_argument("--json", default=None, help="把映射写成 JSON")
    args = ap.parse_args()

    prog = tp.load(args.path)
    reg = dict(prog.module_registry(0))
    _mid, _entry, seq = tdc_layout.find_manifest_module(prog)
    rets: List = load_items(args.items) if args.items else []

    print(f"文件            : {args.path}")
    print(f"清单模块        : registry[{_mid}]")
    print(f"清单项数        : {len(seq)}")
    if args.items:
        print(f"运行时取样      : {len(rets)} 项")
        print(f"推算 cd 槽数    : {total_slots(seq, rets)}")

    if rets:
        print("\n项  模块  工厂    取到的值(截断)                       占槽  首槽")
        for i, m in enumerate(seq):
            e = reg.get(m)
            r = rets[i] if i < len(rets) else None
            s = str(r)[:36]
            n = slots_of_item(r)
            start = slot_map(seq, rets)[i][2]
            print(f"{i:2d}  {m:3d}  @{e:<6} {s:<38} {n:^3}  {start:2d}")

    if args.plain:
        cd = json.loads(open(args.plain, encoding="latin1").read()
                        .rstrip("\x00").strip()).get("cd") or []
        print(f"\n明文实际 cd 槽数: {len(cd)}")
        if rets:
            ok = len(cd) == total_slots(seq, rets)
            print(f">>> 推算 == 实际 ? {ok}")
            if not ok:
                print("    (差 %d 槽: 检查是否有项返回值的逗号计数没算准)" % (len(cd) - total_slots(seq, rets)))
            else:
                print("\n逐项 -> 槽对位:")
                sm = slot_map(seq, rets)
                for (i, m, st) in sm:
                    n = slots_of_item(rets[i])
                    vals = [repr(cd[st + k])[:24] for k in range(min(n, len(cd) - st))]
                    print(f"   项{i:2d}(模块{m:2d}) -> 槽 {st}..{st + n - 1}  {' | '.join(vals)}")

    if args.json and rets:
        Path(args.json).write_text(json.dumps(
            {"manifestModule": _mid, "seq": list(seq),
             "map": [{"item": i, "module": m, "firstSlot": st, "slots": slots_of_item(rets[i])}
                     for (i, m, st) in slot_map(seq, rets)]},
            ensure_ascii=False, indent=1), encoding="utf-8")
        print(f"映射 -> {args.json}")
    return 0


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