"""wupData 信封构造与补丁工具。

解析与改写 WUP 信封中的 cert、uid、session 等字段。
"""
from __future__ import annotations

import base64
import json
import struct
from pathlib import Path

INT8, INT16, INT32, INT64 = 0x00, 0x01, 0x02, 0x03
STRING1, STRING4 = 0x06, 0x07
MAP, LIST = 0x08, 0x09
STRUCT_BEGIN, STRUCT_END = 0x0A, 0x0B
ZERO, SIMPLE_LIST = 0x0C, 0x0D

DEFAULT_QURL_B64 = (
    "AAAD5hADLDxCAFpBBVYMaHV5YXVkYndlYnVpZgdkZWZhdWx0fQABA7gIAAIGCV93dXBfZGF0YR0AAQOKCgoMFgMxLjAm"
    "ynsiYXNzb2NpYXRpb25JZCI6MTg0NTQ5MzkyLCJmdW5jTmFtZSI6IiIsImdyb3VwIjowLCJpZCI6MTg0NTQ5MzkyLCJz"
    "ZXNzaW9uIjo1OTE0ODg1LCJzdGVwIjowLCJzdGlsbExvZ2luIjpmYWxzZSwidHJhY2VJZCI6IjBiOGYwOThmZjY0YTVi"
    "ZGMtMjY2OTItODI2MzkzOTQ3ODc2NjM1NTEzNjUiLCJ0eXBlIjoyLCJ1aWQiOjAsInVzZXJDb250ZXh0IjoiIn02BDUw"
    "MDhAA1a0UFF3ZW1BTjlOSGtaS29NcVdNUW5WUkl5cHFNVGFRRU9ybVhyMzd4UVZoUVpxclA1aVVLRVExMXh2RTB2cE1n"
    "a2xlajIzbEpGYmFHVW5LUEFhYnhWaUF0TnZyTkxBbXhBQzJyRGp6Qy9JSU0vSWFQeTdTcytvQXZqSmltR2pBS2RnVmpr"
    "bUhWV2Q2bEw4cUZScU4zWkJMamt4c2xUQjczaXNvallWcjlrSFhWdUh3SkxoM3YzZgB2AIYAlgALGgYgZWQwZGI4MzM0"
    "Y2FkZDIzNmMwMGNhZGY3ZTExYWI1YTUWBzEzLjQuMjImCTEuMC44MDEzODYARgkxMjcuMC4wLjFWBnhpYW9taWYACyoA"
    "ARYJTTIxMDJKMlNDJigwMmRmMzk4Nzk3NDMyZWFkZWZjYzEyNzY3MTE5YWQ1ZTgwOTk5Mzg5NgdhbmRyb2lkRg9NMjEw"
    "MkoyU0MsMzAsMTFmBDEwODB2BDIxMjCGKDdjNTM4N2UwNTM5YzAyM2MzMWM0ZmYwZTgwN2U3MjU2MTE3Mzg1ZWULMwAA"
    "ARdRuGVvRwAAAQREQ0JHY0QyN0liSWEvZnZ0UHhNT2xZdGJUY0M3bWRaUlJ3YzIyc2NnQkFyRTJ5eTZwSUdIQjNMK0tr"
    "MzVxVS9iaWc1Qk1TVVVSd3gzeS9wWVpWajRjd20rWEg1dnNrakR2SzNhUlhudEJGcURDQUUvaUVIbGs4ZXJ3VUJZdmJM"
    "aXlIb0YrSytQam5GTGJRMmlzSHVhcUtqTHAvWmRETDlxSit3VEVSb3h0ZjFuUzlTZ0l2N3lCaVIyMjd4N3F3RjllUTVu"
    "ckNaRitnRnczelVZb2N6Uk9jbHE1aXZDclhRZTVSZ3hOYkp5aWQ3ZkZqTVhYYlNQdHBIZ2p0TVJtdWp6RVRvPVYAZgAL"
    "BhB3dXB1ZGJyZXF1ZXN0X3YwHQAABQIAWkEFjJgMqAw="
)


def _read_len_int(d: bytes, p: int) -> tuple[int, int]:
    dt = d[p] & 0x0F
    if dt == ZERO:
        return 0, p + 1
    if dt == INT8:
        return struct.unpack_from(">b", d, p + 1)[0], p + 2
    if dt == INT16:
        return struct.unpack_from(">h", d, p + 1)[0], p + 3
    if dt == INT32:
        return struct.unpack_from(">i", d, p + 1)[0], p + 5
    raise ValueError(f"长度int类型异常 {dt:#x}@{p}")


def _skip_value(d: bytes, p: int, dt: int) -> int:
    if dt == ZERO:
        return p
    if dt == INT8:
        return p + 1
    if dt == INT16:
        return p + 2
    if dt == INT32:
        return p + 4
    if dt == INT64:
        return p + 8
    if dt == STRING1:
        return p + 1 + d[p]
    if dt == STRING4:
        return p + 4 + struct.unpack_from(">i", d, p)[0]
    if dt == SIMPLE_LIST:
        p += 1
        n, p = _read_len_int(d, p)
        return p + n
    if dt == MAP:
        n, p = _read_len_int(d, p)
        for _ in range(n):
            h = d[p]
            p += 1
            kd = h & 0x0F
            if kd == STRING1:
                p += 1 + d[p]
            elif kd == STRING4:
                p += 4 + struct.unpack_from(">i", d, p)[0]
            else:
                raise ValueError(f"map key 类型 {kd:#x}")
            vh = d[p]
            p += 1
            p = _skip_value(d, p, vh & 0x0F)
        return p
    if dt == LIST:
        n, p = _read_len_int(d, p)
        eh = d[p]
        p += 1
        edt = eh & 0x0F
        for _ in range(n):
            p = _skip_value(d, p, edt)
        return p
    if dt == STRUCT_BEGIN:
        while True:
            h = d[p]
            p += 1
            sdt = h & 0x0F
            if sdt == STRUCT_END:
                break
            p = _skip_value(d, p, sdt)
        return p
    raise ValueError(f"未知类型 {dt:#x}@{p}")


class Envelope:
    """WUP 请求信封结构解析器与补丁器。"""

    def __init__(self, raw_bytes: bytes):
        self.raw = bytearray(raw_bytes)
        self.tag4_span: tuple[int, int] | None = None
        self.meta_json_span: tuple[int, int] | None = None
        self.uid_off: int | None = None
        self.cert_off: int | None = None
        self.cert_len: int = 260
        self._parse()

    @classmethod
    def load(cls, path: str | Path | None = None) -> "Envelope":
        """加载信封模板，支持从文件加载或使用内嵌金样本。"""
        if path:
            p = Path(path)
            if p.exists():
                return cls._load_from_path(p)
        # 尝试查找 evidence/cert_keycap.json
        candidate = Path(__file__).resolve().parent.parent.parent / "evidence" / "cert_keycap.json"
        if candidate.exists():
            try:
                return cls._load_from_path(candidate)
            except Exception:
                pass
        return cls(base64.b64decode(DEFAULT_QURL_B64))

    @classmethod
    def _load_from_path(cls, p: Path) -> "Envelope":
        if p.suffix == ".json":
            j = json.loads(p.read_text("utf-8"))
            q = next(e["data"] for e in j if e.get("type") == "qurl_done" and e.get("data"))
            return cls(base64.b64decode(q))
        return cls(p.read_bytes())

    def _parse(self) -> None:
        d = self.raw
        p = 4
        # WUP header
        while p < len(d):
            h = d[p]
            p += 1
            tag, dt = (h >> 4) & 0x0F, h & 0x0F
            if tag == 4 and dt == INT32:
                self.tag4_span = (p, p + 4)
                p += 4
            elif tag == 7 and dt == SIMPLE_LIST:
                p += 1  # 元素类型头
                n, p = _read_len_int(d, p)
                self._parse_sbuffer(p, n)
                break
            else:
                p = _skip_value(d, p, dt)

    def _parse_sbuffer(self, start: int, ln: int) -> None:
        d = self.raw
        p = start
        h = d[p]
        p += 1
        assert (h & 0x0F) == MAP
        cnt, p = _read_len_int(d, p)
        for _ in range(cnt):
            kh = d[p]
            p += 1
            kln = d[p]
            p += 1
            k = bytes(d[p:p + kln])
            p += kln
            vh = d[p]
            p += 1
            if k == b"_wup_data":
                assert (vh & 0x0F) == SIMPLE_LIST
                p += 1
                wup_data_len, p = _read_len_int(d, p)
                self._parse_wup_data_struct(p, wup_data_len)
                p += wup_data_len
            else:
                p = _skip_value(d, p, vh & 0x0F)

    def _parse_wup_data_struct(self, start: int, ln: int) -> None:
        d = self.raw
        q = start
        assert (d[q] & 0x0F) == STRUCT_BEGIN
        q += 1
        while q < start + ln:
            hh = d[q]
            q += 1
            tag, dt = (hh >> 4) & 0x0F, hh & 0x0F
            if tag == 15:
                tag = d[q]
                q += 1
            if dt == STRUCT_END:
                break
            if tag == 0 and dt == STRUCT_BEGIN:
                while True:
                    h2 = d[q]
                    q += 1
                    tg2, dt2 = (h2 >> 4) & 0x0F, h2 & 0x0F
                    if tg2 == 15:
                        tg2 = d[q]
                        q += 1
                    if dt2 == STRUCT_END:
                        break
                    vs = q
                    q = _skip_value(d, q, dt2)
                    if tg2 == 2 and dt2 in (STRING1, STRING4):
                        off = vs + (4 if dt2 == STRING4 else 1)
                        self.meta_json_span = (off, q)
            elif tag == 3 and dt == INT64:
                self.uid_off = q
                q += 8
            elif tag == 4 and dt == STRING4:
                ln_c = struct.unpack_from(">i", d, q)[0]
                self.cert_off = q + 4
                self.cert_len = ln_c
                q = self.cert_off + ln_c
            else:
                q = _skip_value(d, q, dt)
        if self.uid_off is None or self.cert_off is None:
            raise ValueError("未在信封中定位到 t3(uid) 或 t4(cert)")

    @property
    def uid(self) -> int:
        return struct.unpack_from(">Q", self.raw, self.uid_off)[0]

    def patch_uid(self, uid: int) -> "Envelope":
        struct.pack_into(">Q", self.raw, self.uid_off, uid)
        return self

    @property
    def cert_b64(self) -> bytes:
        return bytes(self.raw[self.cert_off:self.cert_off + self.cert_len])

    def patch_cert(self, cert: bytes) -> "Envelope":
        b64 = base64.b64encode(cert)
        if len(b64) != self.cert_len:
            raise ValueError(
                f"证书b64长度不符: {len(b64)} != 模板 {self.cert_len} (cert {len(cert)}B)"
            )
        self.raw[self.cert_off:self.cert_off + self.cert_len] = b64
        return self

    def patch_session(self, session: int) -> "Envelope":
        if self.tag4_span:
            struct.pack_into(">I", self.raw, self.tag4_span[0], session & 0xFFFFFFFF)
        d = self.raw
        i = d.find(b"wupudbrequest_v0")
        if i >= 0:
            j = i + len(b"wupudbrequest_v0") + 4
            if j + 4 <= len(d):
                struct.pack_into(">I", d, j, session & 0xFFFFFFFF)
        return self

    def wup_b64(self) -> str:
        return base64.b64encode(bytes(self.raw)).decode()