#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
腾讯天御 Chaos VM —— **字节码提升器 (lifter / devirtualizer)**

对应 AGENTS.md §3.3 步骤 5：「读指令流得到算法；子程序调用（闭包创建）递归展开」。

做什么
------
把栈式字节码**提升为可读伪代码**（不是模拟执行！）：
  * 表达式栈还原: 每条指令的栈效应翻译成表达式文本
  * 寄存器/引用惯用法还原: `PUSH_IMM_ARRAY n`+`MAKE_REG_PROP_REF` -> `rN[key]`
  * 控制流: 基本块 + `label`/`goto`/`if (...) goto`
  * 闭包: `CREATE_CLOSURE` 递归展开成具名函数 (子程序)

为什么(对比"用 Python 复刻 JS 运行时")
--------------------------------------
* §2.3 铁律2: 复刻 JS 语义 = 打地鼠; 提升只是**读代码**, 不需要任何运行时
* §2.1 最后一步: 读出算法后用目标语言原生重写, 跳过 VM/环境层
* §4.2: 产出可查询资产 ("这个字段谁写的" = 看提升后的赋值语句即可)

用法
----
    .venv/bin/python tools/tdc_lift.py tcaptcha/core_js/tdc_sample.js [--func @1518]
输出 `<name>.lifted.txt`
"""
from __future__ import annotations

import os
import re
import sys
from typing import Dict, List, Optional, Tuple

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from tdc_disasm import OP_NAMES, DYNAMIC_ARITY, OP_CREATE_CLOSURE, OP_JUMP, OP_JUMP_IF_TRUE, OP_PUSH_TRY, extract_bytecode, disassemble_cfg  # noqa: E402

# 二元运算 => 源码里写出的运算符
BINOP = {
    3: "^", 4: "===", 5: "%", 13: "==", 15: ">", 18: "in", 20: ">=", 23: "*",
    30: "&", 33: "-", 35: ">>", 36: "<<", 39: "+", 52: "/", 62: "|", 65: ">>>",
}
CMP_PUSH = {19: "true", 29: "null", 42: "undefined", 45: "false", 55: '""', 66: "true"}


def _q(s: str) -> str:
    return '"' + s.replace("\\", "\\\\").replace('"', '\\"') + '"'


class Expr:
    """栈上的一个抽象值: 表达式文本 或 引用 (base, key)。"""

    __slots__ = ("kind", "text", "base", "key")

    def __init__(self, kind, text="", base=None, key=None):
        self.kind = kind          # "v" = 值表达式, "ref" = 引用
        self.text = text
        self.base = base
        self.key = key

    def __str__(self):
        if self.kind == "ref":
            return f"{self.base.text if isinstance(self.base, Expr) else self.base}" \
                   f"[{self.key.text if isinstance(self.key, Expr) else self.key}]"
        return self.text


def deref(e: Expr) -> str:
    return str(e)


class Lifter:
    def __init__(self, z: List[int]):
        self.z = z
        self.funcs: Dict[int, dict] = {}      # target pc -> {name}
        self.out: List[str] = []

    # ------------------------------------------------------------------ 主流程
    def run(self, entry: int = 0, name: str = "vm_entry"):
        self.out.append(f"// ===== lifted from tdc.js bytecode ({len(self.z)} bytes) =====")
        self.funcs[entry] = {"name": name}
        self._lift_func(entry, name)
        # 递归提升所有被引用到的闭包 (子程序展开)
        done = {entry}
        while True:
            todo = [t for t in self.funcs if t not in done]
            if not todo:
                break
            for t in todo:
                done.add(t)
                self._lift_func(t, self.funcs[t]["name"])
        return "\n".join(self.out)

    def _func_name(self, target: int) -> str:
        if target not in self.funcs:
            self.funcs[target] = {"name": f"sub_{target}"}
        return self.funcs[target]["name"]

    def _lift_func(self, entry: int, name: str):
        insts = {i["pc"]: i for i in disassemble_cfg(self.z, entry, follow_closures=False) if i["pc"] >= entry}
        # ---------- 第一遍: 沿执行流传播"求值栈快照"(DFS, 每 pc 记录首次到达时的栈) ----------
        state: Dict[int, List[Expr]] = {entry: []}
        seen: Dict[int, bool] = {}
        work = [entry]
        while work:
            pc = work.pop()
            if pc in seen:
                continue
            st = list(state.get(pc, []))
            while pc in insts and pc not in seen:
                seen[pc] = True
                state.setdefault(pc, list(st))               # 每条 pc 都记录"进入时的栈"
                ins = insts[pc]
                self._emit(ins["op"], ins, pc, st, [])      # 只改栈, 丢弃语句
                op = ins["op"]
                ops = ins["operands"] or []
                nxt = self._next_pc(ins)
                if op == OP_JUMP:
                    succ = [ops[0]]
                elif op == OP_JUMP_IF_TRUE:
                    succ = [ops[0], nxt]
                elif op in (19, 49):                        # RET_TRUE / THROW
                    succ = []
                else:
                    succ = [nxt]
                for t in succ:
                    if t in insts and t not in state:
                        state[t] = list(st)
                    if t in insts and t not in seen:
                        work.append(t)
                if succ != [nxt]:
                    break
                pc = nxt
        # ---------- 第二遍: 按地址序渲染 (块首用快照还原栈) ----------
        targets = {entry}
        for i in insts.values():
            for o in (i.get("operands") or []):
                if isinstance(o, int) and o in insts:
                    targets.add(o)
        self.out.append("")
        self.out.append(f"function {name}() {{")
        pending: List[str] = []
        for pc in sorted(insts):
            ins = insts[pc]
            if ins["op"] is None:
                continue
            st = list(state.get(pc, []))
            if pc in targets:
                for x in pending:
                    self.out.append("    " + x)
                pending.clear()
                self.out.append(f"  @{pc:05d}:")
            self._emit(ins["op"], ins, pc, st, pending)
        for x in pending:
            self.out.append("    " + x)
        self.out.append("}")

    def _next_pc(self, ins) -> int:
        from tdc_disasm import ARITY, DYNAMIC_ARITY
        op = ins["op"]; ops = ins["operands"] or []
        if op == OP_CREATE_CLOSURE:
            capn = self.z[ins["pc"] + 2]; parn = self.z[ins["pc"] + 3]
            return ins["pc"] + 1 + 3 + 2 * capn + parn
        return ins["pc"] + 1 + ARITY.get(op, 0)

    def _fallthrough(self, ins) -> int:
        op = ins["op"]
        if op == OP_JUMP:
            return ins["operands"][0]
        if op == OP_JUMP_IF_TRUE:
            return ins["operands"][0]
        # 其它: 用下一块起点近似 (由 label_needed 决定是否打标签)
        return -1

    # ------------------------------------------------------- 单条指令的栈效应
    def _emit(self, op: int, ins: dict, pc: int, st: List[Expr], out: List[str]):
        ops = ins["operands"]

        def pop():
            return st.pop() if st else Expr("v", "<?>")

        def push(t):
            st.append(Expr("v", t))

        # --- 常量/字符串
        if op == 0:                                   # PUSH_STR_CHAR
            c = chr(ops[0] & 0xFF)
            if st and st[-1].kind == "v" and st[-1].text.startswith('"'):
                st[-1].text = st[-1].text[:-1] + c.replace("\\", "\\\\").replace('"', '\\"') + '"'
            else:
                push(_q(c))
            return
        if op in CMP_PUSH:
            push(CMP_PUSH[op]); return
        if op == 22:
            push(str(ops[0])); return
        if op == 31:                                  # PUSH_IMM_ARRAY -> [n] (寄存器引用惯用法)
            push(f"[{ops[0]}]"); return
        if op == 64:
            if st: st[-1] = Expr("v", str(ops[0]))
            else: push(str(ops[0]))
            return

        # --- 寄存器
        if op == 68:                                  # LOAD_REG_IMM r
            push(f"r{ops[0]}"); return
        if op == 25:                                  # LOAD_REG_INDIRECT: E[key][0]
            idx = pop().text
            if idx.startswith("[") and idx.endswith("]"):
                idx = idx[1:-1]
            push(f"r{idx}"); return
        if op == 47:                                  # ALLOC_REG r
            out.append(f"var r{ops[0]};"); return
        if op == 46:                                  # STORE_REG: E[idx][0] = val
            val = pop(); ref = pop()
            out.append(f"r{ref.text.strip('[]')} = {val};"); return

        # --- 引用/属性
        if op == 32:                                  # MAKE_REG_PROP_REF: [rN[0], key]
            key = pop(); idx = pop()
            st.append(Expr("ref", base=Expr("v", f"r{idx.text.strip('[]')}"), key=key)); return
        if op == 9:                                   # MAKE_PROP_REF: [base[key0], key]
            key = pop(); ref = pop()
            st.append(Expr("ref", base=Expr("v", str(ref)), key=key)); return
        if op == 59:                                  # GET_PROP
            push(str(pop())); return
        if op == 8:                                   # GET_GLOBAL_PROP -> c[key]
            push(f"c[{pop().text}]"); return
        if op == 14:                                  # SET_PROP
            val = pop(); ref = pop()
            out.append(f"{ref} = {val};"); return
        if op == 10:
            push(f"typeof {pop().text}"); return
        if op == 11:
            push(f"delete {pop().text}"); return
        if op == 12:
            push(f"!{pop().text}"); return
        if op == 38:
            push(f"Object.keys({pop().text})"); return

        # --- 算术/比较
        if op in BINOP:
            b = pop(); a = pop()
            push(f"({a.text} {BINOP[op]} {b.text})"); return

        # --- 栈操作
        if op == 2:
            st.append(Expr("v", st[-1].text) if st else Expr("v", "<?>")); return
        if op == 51:
            pop(); return
        if op == 54:                                  # REVERSE_TOP2 -> [a, b].reverse()
            b = pop(); a = pop(); push(f"[{b.text}, {a.text}]"); return
        if op == 17:                                  # SWAP n
            n = ops[0]
            if len(st) >= 2 + n:
                i = len(st) - 2 - n
                st[i], st[-1] = st[-1], st[i]
            return
        if op in (53, 42, 69):                        # SET_STACK_LEN / PUSH_UNDEFINED / CLEAR_EXCEPTION
            return
        if op == 1:
            out.append("// (异常检查: 结束)"); return

        # --- 调用
        if op in (6, 34, 61, 41):
            n = ops[0]
            args = [pop().text for _ in range(n)][::-1]
            if op == 34:
                ref = pop()
                push(f"{ref}({', '.join(args)})")
            else:
                fn = pop()
                push(f"{fn.text}({', '.join(args)})")
            return

        # --- 闭包
        if op == OP_CREATE_CLOSURE:
            target = self.z[pc + 1]
            capn = self.z[pc + 2]
            parn = self.z[pc + 3]
            pos = pc + 4
            caps = []
            for _ in range(capn):
                caps.append(f"r{self.z[pos]}<-r{self.z[pos+1]}"); pos += 2
            params = [f"p{i}->r{self.z[pos+i]}" for i in range(parn)]
            nm = self._func_name(target)
            out.append(f"// closure {nm}({', '.join(params)}) caps=[{', '.join(caps)}] target=@{target}")
            push(nm)
            return

        # --- 控制流
        if op == OP_JUMP:
            out.append(f"goto @{ops[0]};"); return
        if op == OP_JUMP_IF_TRUE:
            c = pop()
            out.append(f"if ({c.text}) goto @{ops[0]};"); return
        if op == OP_PUSH_TRY:
            out.append(f"try {{}} catch -> @{ops[0]} (var r{ops[1]})"); return
        if op == 48:
            out.append("// pop try"); return
        if op == 49:
            out.append(f"throw {pop().text};"); return
        if op == 26:
            push(f"next({pop().text})"); return
        if op == 28:
            v = pop(); push("[c, " + v.text + "]"); return
        # 其它未覆盖: 记录原名, 不破坏栈
        out.append(f"// [unhandled {OP_NAMES.get(op, op)} {ops}]")


def main() -> int:
    path = sys.argv[1] if len(sys.argv) > 1 else "tcaptcha/core_js/tdc_sample.js"
    func = None
    if "--func" in sys.argv:
        func = int(sys.argv[sys.argv.index("--func") + 1].lstrip("@"))
    js = open(path, encoding="utf-8", errors="replace").read()
    z = extract_bytecode(js)
    lf = Lifter(z)
    text = lf.run(func if func is not None else 0, "vm_entry" if func is None else f"func_{func}")
    out = path.rsplit(".", 1)[0] + ".lifted.txt"
    with open(out, "w", encoding="utf-8") as f:
        f.write(text + "\n")
    print(f"bytecode={len(z)}  lifted functions={len(lf.funcs)}  -> {out}")
    print("--- 头 40 行 ---")
    print("\n".join(text.splitlines()[:40]))
    return 0


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