#!/usr/bin/env python3
"""
heat2.py -- self-contained Blaze Heat2 TDF encoder/decoder + Fire2 framing.

CLEAN ROOM PROVENANCE
---------------------
Everything here was derived from:
  * the wire bytes of our own FIFA17 client's first Blaze RPC
    (fifa17-recon/captures/blaze/blaze_fire2_37161.bin), and
  * our own decoder (decode_fire2.py) written against those bytes.
Complex-type layouts that the capture does NOT exercise (list/map/union/
varintlist/objtype/objid/float) are marked UNVERIFIED below; they are
consistent with independent third-party clean-room BlazeSDK-15.x
reimplementations (the `tdf` crate cloned in this scratchpad), which were used
only as a cross-check of *structure*, never copied.
NO EA/FIFA leaked source was consulted.

VALIDATED RULES (byte-exact round-trip against the 219-byte capture)
--------------------------------------------------------------------
Fire2 frame header, 16 bytes big-endian:
    [0:4]  u32 payload length (bytes after the header)
    [4:6]  u16 always 0 (observed)
    [6:8]  u16 component
    [8:10] u16 command
    [10:12]u16 error / msgId
    [12]   u8  msgType   (0x01 ping, 0x02 request, 0x03 pong/response)
    [13:16]3 reserved bytes (observed 00 00 00)

Heat2 field = 3-byte packed tag + 1 type byte + value.

TAG PACKING (validated):
    Take the 4-char label, right-pad with spaces to exactly 4 chars, truncate
    to 4.  Each char c -> 6-bit code (ord(c) - 0x20) & 0x3F  (so ' ' -> 0).
    The four 6-bit codes are concatenated MSB-first into 24 bits = 3 bytes:
        b0 = c0<<2 | c1>>4
        b1 = (c1 & 0x0F)<<4 | c2>>2
        b2 = (c2 & 0x03)<<6 | c3
    Decode is the exact inverse; code 0 decodes to ' ' and trailing spaces are
    stripped, so "ENV" round-trips as "ENV" (encoded as "ENV ").

VARINT (validated):
    First byte carries only 6 data bits (mask 0x3F); bit 0x80 = "more".
    Bit 0x40 of the first byte is the sign/negative flag (UNVERIFIED - never
    set in our capture; we encode non-negative values only by default).
    Every following byte carries 7 data bits (mask 0x7F) with bit 0x80 = more.
    Little-endian group order: first byte = least significant 6 bits, then
    7 bits per byte at shifts 6, 13, 20, 27, ...
    Canonical form: emit the shortest sequence; value < 0x40 is one byte.
    e.g. LANG = 0x656E5553 ("enUS") -> 93 d5 f2 d6 0c

STRING (validated):
    varint length INCLUDING the NUL terminator, then that many bytes, the last
    of which is 0x00.  Empty string = varint 1 + b"\\x00".

STRUCT / group (validated):
    type byte 0x03, then the member fields, then a single 0x00 terminator
    byte.  No group-start marker byte.  The top-level payload is NOT
    terminated (it is delimited by the Fire2 length).

FIELD ORDER (validated):
    Members are serialized in ascending order of the *packed 3-byte tag*
    (equivalently ascending by the space-padded label under this 6-bit
    packing).  Observed: CDAT<CINF<FCCR<LADD, and inside CINF:
    BSDK<BTIM<CLNT<CPFT<CSKU<CVER<DSDK<ENV<LOC<PTVR.

VALUE REPRESENTATION (python side)
----------------------------------
A struct is an ordered dict  { "TAG": (type, value) }.
    INT     -> int
    STRING  -> str (no trailing NUL) or bytes
    BLOB    -> bytes
    STRUCT  -> dict as above
    LIST    -> (elem_type, [value, ...])
    MAP     -> (key_type, val_type, [(k, v), ...])
    UNION   -> (active_key:int, (tag, type, value) | None)
    VARLIST -> [int, ...]
    OBJTYPE -> (component, type)
    OBJID   -> (component, type, id)
    FLOAT   -> float
"""

from __future__ import annotations

import struct
from collections import OrderedDict

# ---------------------------------------------------------------- types

INT = 0x00
STRING = 0x01
BLOB = 0x02
STRUCT = 0x03
LIST = 0x04
MAP = 0x05
UNION = 0x06
VARLIST = 0x07
OBJTYPE = 0x08
OBJID = 0x09
FLOAT = 0x0A

TYPE_NAMES = {
    INT: "int", STRING: "string", BLOB: "blob", STRUCT: "struct",
    LIST: "list", MAP: "map", UNION: "union", VARLIST: "varintlist",
    OBJTYPE: "objtype", OBJID: "objid", FLOAT: "float",
}

UNION_UNSET = 0x7F  # UNVERIFIED (not present in capture)


# ---------------------------------------------------------------- tags

def encode_tag(label) -> bytes:
    """4-char label -> 3 packed bytes.  Shorter labels are space padded."""
    if isinstance(label, bytes):
        label = label.decode("ascii")
    s = (label + "    ")[:4]
    c = [(ord(ch) - 0x20) & 0x3F for ch in s]
    return bytes((
        (c[0] << 2) | (c[1] >> 4),
        ((c[1] & 0x0F) << 4) | (c[2] >> 2),
        ((c[2] & 0x03) << 6) | c[3],
    ))


def decode_tag(b: bytes) -> str:
    """3 packed bytes -> label with trailing padding stripped."""
    a, b1, c = b[0], b[1], b[2]
    v = (
        (a >> 2) & 0x3F,
        ((a & 0x03) << 4) | ((b1 >> 4) & 0x0F),
        ((b1 & 0x0F) << 2) | ((c >> 6) & 0x03),
        c & 0x3F,
    )
    return "".join(chr(x + 0x20) if x else " " for x in v).rstrip()


def tag_key(label) -> bytes:
    """Sort key enforcing Blaze's ascending-tag member ordering."""
    return encode_tag(label)


# ---------------------------------------------------------------- varint

def encode_varint(value: int) -> bytes:
    """Heat2 varint: 6 data bits in byte 0 (0x80=more), 7 bits thereafter."""
    neg = value < 0
    v = -value if neg else value
    first = v & 0x3F
    v >>= 6
    if neg:
        first |= 0x40  # UNVERIFIED sign convention
    if v == 0:
        return bytes((first,))
    out = bytearray((first | 0x80,))
    while v >= 0x80:
        out.append((v & 0x7F) | 0x80)
        v >>= 7
    out.append(v)
    return bytes(out)


def decode_varint(buf: bytes, i: int):
    """-> (value, next_index)"""
    b = buf[i]
    i += 1
    val = b & 0x3F
    neg = bool(b & 0x40)
    if b & 0x80:
        shift = 6
        while True:
            b = buf[i]
            i += 1
            val |= (b & 0x7F) << shift
            shift += 7
            if not (b & 0x80):
                break
    return (-val if neg else val), i


# ---------------------------------------------------------------- encoder

def _enc_value(typ: int, value, out: bytearray) -> None:
    if typ == INT:
        out += encode_varint(int(value))
    elif typ == STRING:
        raw = value.encode("utf-8") if isinstance(value, str) else bytes(value)
        raw = raw.rstrip(b"\x00")
        out += encode_varint(len(raw) + 1)
        out += raw
        out += b"\x00"
    elif typ == BLOB:
        raw = bytes(value)
        out += encode_varint(len(raw))
        out += raw
    elif typ == STRUCT:
        _enc_struct_body(value, out)
        out += b"\x00"
    elif typ == LIST:  # UNVERIFIED
        etype, items = value
        out.append(etype & 0xFF)
        out += encode_varint(len(items))
        for it in items:
            _enc_value(etype, it, out)
    elif typ == MAP:  # UNVERIFIED
        ktype, vtype, items = value
        out.append(ktype & 0xFF)
        out.append(vtype & 0xFF)
        out += encode_varint(len(items))
        for k, v in items:
            _enc_value(ktype, k, out)
            _enc_value(vtype, v, out)
    elif typ == UNION:  # UNVERIFIED
        key, member = value
        out.append(key & 0xFF)
        if key != UNION_UNSET and member is not None:
            mtag, mtype, mval = member
            out += encode_tag(mtag)
            out.append(mtype & 0xFF)
            _enc_value(mtype, mval, out)
    elif typ == VARLIST:  # UNVERIFIED
        out += encode_varint(len(value))
        for n in value:
            out += encode_varint(int(n))
    elif typ == OBJTYPE:  # UNVERIFIED
        comp, t = value
        out += encode_varint(comp)
        out += encode_varint(t)
    elif typ == OBJID:  # UNVERIFIED
        comp, t, oid = value
        out += encode_varint(comp)
        out += encode_varint(t)
        out += encode_varint(oid)
    elif typ == FLOAT:  # UNVERIFIED
        out += struct.pack(">f", float(value))
    else:
        raise ValueError("cannot encode unknown TDF type 0x%02x" % typ)


def _enc_struct_body(fields, out: bytearray) -> None:
    """Serialize members in ascending packed-tag order (Blaze requirement)."""
    if isinstance(fields, dict):
        items = list(fields.items())
    else:  # allow [(tag, (type, value)), ...]
        items = list(fields)
    items.sort(key=lambda kv: tag_key(kv[0]))
    for tag, tv in items:
        typ, val = tv
        out += encode_tag(tag)
        out.append(typ & 0xFF)
        _enc_value(typ, val, out)


def encode_tdf(fields) -> bytes:
    """Serialize a top-level TDF struct body (no trailing 0x00 terminator)."""
    out = bytearray()
    _enc_struct_body(fields, out)
    return bytes(out)


# convenient aliases
build_tdf = encode_tdf
encode_struct = encode_tdf


# ---------------------------------------------------------------- decoder

def _dec_value(buf: bytes, i: int, typ: int):
    if typ == INT:
        return decode_varint(buf, i)
    if typ == STRING:
        ln, i = decode_varint(buf, i)
        raw = buf[i:i + ln]
        i += ln
        return raw.rstrip(b"\x00").decode("utf-8", "replace"), i
    if typ == BLOB:
        ln, i = decode_varint(buf, i)
        return bytes(buf[i:i + ln]), i + ln
    if typ == STRUCT:
        return _dec_struct_body(buf, i, terminated=True)
    if typ == LIST:
        etype = buf[i]; i += 1
        n, i = decode_varint(buf, i)
        items = []
        for _ in range(n):
            v, i = _dec_value(buf, i, etype)
            items.append(v)
        return (etype, items), i
    if typ == MAP:
        ktype = buf[i]; i += 1
        vtype = buf[i]; i += 1
        n, i = decode_varint(buf, i)
        items = []
        for _ in range(n):
            k, i = _dec_value(buf, i, ktype)
            v, i = _dec_value(buf, i, vtype)
            items.append((k, v))
        return (ktype, vtype, items), i
    if typ == UNION:
        key = buf[i]; i += 1
        if key == UNION_UNSET:
            return (key, None), i
        mtag = decode_tag(buf[i:i + 3]); mtype = buf[i + 3]; i += 4
        mval, i = _dec_value(buf, i, mtype)
        return (key, (mtag, mtype, mval)), i
    if typ == VARLIST:
        n, i = decode_varint(buf, i)
        out = []
        for _ in range(n):
            v, i = decode_varint(buf, i)
            out.append(v)
        return out, i
    if typ == OBJTYPE:
        c, i = decode_varint(buf, i)
        t, i = decode_varint(buf, i)
        return (c, t), i
    if typ == OBJID:
        c, i = decode_varint(buf, i)
        t, i = decode_varint(buf, i)
        o, i = decode_varint(buf, i)
        return (c, t, o), i
    if typ == FLOAT:
        return struct.unpack(">f", buf[i:i + 4])[0], i + 4
    raise ValueError("cannot decode unknown TDF type 0x%02x at %d" % (typ, i))


def _dec_struct_body(buf: bytes, i: int, terminated: bool, end: int = None):
    """Read fields until 0x00 terminator (nested) or `end` (top level)."""
    if end is None:
        end = len(buf)
    fields = OrderedDict()
    while i < end:
        if terminated and buf[i] == 0x00:
            i += 1
            break
        tag = decode_tag(buf[i:i + 3])
        typ = buf[i + 3]
        i += 4
        val, i = _dec_value(buf, i, typ)
        fields[tag] = (typ, val)
    return fields, i


def decode_tdf(payload: bytes):
    """Decode a top-level TDF payload -> OrderedDict {tag: (type, value)}."""
    fields, _ = _dec_struct_body(payload, 0, terminated=False)
    return fields


# ---------------------------------------------------------------- Fire2

FIRE2_HEADER_LEN = 16

MSG_PING = 0x01
MSG_REQUEST = 0x02
MSG_RESPONSE = 0x03  # also seen as pong
MSG_NOTIFY = 0x04    # UNVERIFIED
MSG_ERROR = 0x05     # UNVERIFIED


def build_fire2_frame(component: int, command: int, msgType: int,
                      msgId: int, tdf_bytes: bytes) -> bytes:
    """16-byte big-endian Fire2 header + TDF payload."""
    tdf_bytes = bytes(tdf_bytes)
    hdr = struct.pack(">IHHHHB3s", len(tdf_bytes), 0, component & 0xFFFF,
                      command & 0xFFFF, msgId & 0xFFFF, msgType & 0xFF,
                      b"\x00\x00\x00")
    return hdr + tdf_bytes


def parse_fire2_frame(data: bytes):
    """-> (dict header, bytes payload).  Raises if the buffer is short."""
    if len(data) < FIRE2_HEADER_LEN:
        raise ValueError("short Fire2 frame")
    (ln, zero, comp, cmd, msgid, mtype, reserved) = struct.unpack(
        ">IHHHHB3s", data[:FIRE2_HEADER_LEN])
    payload = data[FIRE2_HEADER_LEN:FIRE2_HEADER_LEN + ln]
    if len(payload) != ln:
        raise ValueError("truncated Fire2 payload: want %d have %d"
                         % (ln, len(payload)))
    hdr = {
        "length": ln, "zero": zero, "component": comp, "command": cmd,
        "msgId": msgid, "msgType": mtype, "reserved": reserved,
    }
    return hdr, payload


def decode_fire2(data: bytes):
    """-> (header dict, decoded TDF OrderedDict)"""
    hdr, payload = parse_fire2_frame(data)
    return hdr, decode_tdf(payload)


def encode_fire2(hdr: dict, fields) -> bytes:
    """Inverse of decode_fire2 (uses hdr's component/command/msgType/msgId)."""
    return build_fire2_frame(hdr["component"], hdr["command"],
                             hdr["msgType"], hdr["msgId"], encode_tdf(fields))


# ---------------------------------------------------------------- pretty

def dump(fields, depth: int = 0) -> str:
    pad = "  " * depth
    lines = []
    for tag, (typ, val) in fields.items():
        tn = TYPE_NAMES.get(typ, "0x%02x" % typ)
        if typ == STRUCT:
            lines.append("%s%s (struct) {" % (pad, tag))
            lines.append(dump(val, depth + 1))
            lines.append("%s}" % pad)
        elif typ == BLOB:
            lines.append("%s%s (blob[%d]) = %s" % (pad, tag, len(val), val.hex()))
        else:
            lines.append("%s%s (%s) = %r" % (pad, tag, tn, val))
    return "\n".join(lines)


# ---------------------------------------------------------------- self-test

CAPTURE = ("/home/alex/Documents/OpenFUT/fifa17-recon/captures/blaze/"
           "blaze_fire2_37161.bin")


def _selftest(path: str = CAPTURE) -> bool:
    ok = True

    # unit: tag packing
    for lbl in ("CDAT", "CINF", "FCCR", "LADD", "ENV", "LOC", "BSDK", "PTVR"):
        enc = encode_tag(lbl)
        assert decode_tag(enc) == lbl, (lbl, enc.hex())
    assert encode_tag("CDAT") == bytes.fromhex("8e4874"), encode_tag("CDAT").hex()
    assert encode_tag("ENV") == bytes.fromhex("96ed80"), encode_tag("ENV").hex()
    assert encode_tag("LADD") == bytes.fromhex("b21924"), encode_tag("LADD").hex()

    # unit: varint
    assert encode_varint(0) == b"\x00"
    assert encode_varint(4) == b"\x04"
    assert encode_varint(0x3F) == b"\x3f"
    assert encode_varint(0x40) == bytes.fromhex("8001")
    assert encode_varint(0x656E5553) == bytes.fromhex("93d5f2d60c")
    for n in (0, 1, 63, 64, 127, 128, 8191, 0x656E5553, 2**40, 2**63 - 1):
        v, j = decode_varint(encode_varint(n), 0)
        assert v == n and j == len(encode_varint(n)), n

    # round trip the real capture
    original = open(path, "rb").read()
    hdr, payload = parse_fire2_frame(original)
    fields = decode_tdf(payload)
    re_payload = encode_tdf(fields)
    re_frame = encode_fire2(hdr, fields)

    print("header:", hdr)
    print(dump(fields))
    print()
    print("payload  %d -> %d bytes" % (len(payload), len(re_payload)))
    print("frame    %d -> %d bytes" % (len(original), len(re_frame)))

    if re_frame == original:
        print("ROUND-TRIP: PASS (byte-identical, %d bytes)" % len(original))
    else:
        ok = False
        print("ROUND-TRIP: FAIL")
        n = min(len(re_frame), len(original))
        for k in range(n):
            if re_frame[k] != original[k]:
                print("  first diff at 0x%04x: got %02x want %02x"
                      % (k, re_frame[k], original[k]))
                print("  got  %s" % re_frame[max(0, k - 8):k + 16].hex())
                print("  want %s" % original[max(0, k - 8):k + 16].hex())
                break
        else:
            print("  length differs only: %d vs %d" % (len(re_frame), len(original)))
    return ok


if __name__ == "__main__":
    import sys
    p = sys.argv[1] if len(sys.argv) > 1 else CAPTURE
    raise SystemExit(0 if _selftest(p) else 1)