#!/usr/bin/env python3
"""Extract a playable web bundle from a WC3 d3d9 apitrace.

Feed:  apitrace pickle war3.trace | python3 extract.py OUTDIR
Emits: OUTDIR/bundle.json (meta+frames) + OUTDIR/blobs.bin (dedup blob store)

Records RECORD_FRAMES frames once per-frame draw count exceeds ACTIVATE_DRAWS
(i.e. real gameplay, not menus). Buffers snapshot at draw time by content
hash — the dedup ratio printed at the end is the transport-feasibility number.
"""
import hashlib
import json
import os
import struct
import sys

sys.path.insert(0, "/home/ubuntu/apitrace-14.0-Linux/lib/apitrace/scripts")
import unpickle

START_FRAME = int(os.environ.get("START_FRAME", "2600"))
RECORD_FRAMES = int(os.environ.get("RECORD_FRAMES", "240"))

OUT = sys.argv[1] if len(sys.argv) > 1 else "bundle"
os.makedirs(OUT, exist_ok=True)

DRAWS = {"DrawPrimitive", "DrawIndexedPrimitive",
         "DrawPrimitiveUP", "DrawIndexedPrimitiveUP"}
REL_STATES = {7, 14, 15, 19, 20, 22, 23, 24, 25, 26, 27, 28, 29, 34, 35,
              36, 37, 137, 139, 140, 141, 145, 8, 9, 168}


class Blobs:
    def __init__(self):
        self.f = open(os.path.join(OUT, "blobs.bin"), "wb")
        self.index = {}
        self.entries = []
        self.off = 0
        self.raw_bytes = 0        # bytes before dedup (dynamic traffic)
    def put(self, data: bytes) -> int:
        self.raw_bytes += len(data)
        h = hashlib.blake2b(data, digest_size=16).digest()
        if h in self.index:
            return self.index[h]
        i = len(self.entries)
        self.index[h] = i
        self.entries.append((self.off, len(data)))
        self.f.write(data)
        self.off += len(data)
        return i


def decode_texture(fmt, w, h, data, pitch):
    """Return (RGBA8 bytes, w, h) for the formats WC3 actually uses."""
    out = bytearray(w * h * 4)
    if fmt in (21, 22):                      # A8R8G8B8 / X8R8G8B8
        for y in range(h):
            row = data[y*pitch : y*pitch + w*4]
            for x in range(w):
                b, g, r, a = row[x*4:x*4+4]
                i = (y*w + x) * 4
                out[i:i+4] = bytes((r, g, b, 255 if fmt == 22 else a))
    elif fmt == 23:                          # R5G6B5
        for y in range(h):
            for x in range(w):
                v, = struct.unpack_from("<H", data, y*pitch + x*2)
                i = (y*w + x) * 4
                out[i:i+4] = bytes((((v >> 11) & 31) * 255 // 31,
                                    ((v >> 5) & 63) * 255 // 63,
                                    (v & 31) * 255 // 31, 255))
    elif fmt in (24, 25):                    # X1R5G5B5 / A1R5G5B5
        for y in range(h):
            for x in range(w):
                v, = struct.unpack_from("<H", data, y*pitch + x*2)
                i = (y*w + x) * 4
                a = 255 if fmt == 24 or (v & 0x8000) else 0
                out[i:i+4] = bytes((((v >> 10) & 31) * 255 // 31,
                                    ((v >> 5) & 31) * 255 // 31,
                                    (v & 31) * 255 // 31, a))
    elif fmt == 26:                          # A4R4G4B4
        for y in range(h):
            for x in range(w):
                v, = struct.unpack_from("<H", data, y*pitch + x*2)
                i = (y*w + x) * 4
                out[i:i+4] = bytes((((v >> 8) & 15) * 17, ((v >> 4) & 15) * 17,
                                    (v & 15) * 17, ((v >> 12) & 15) * 17))
    elif fmt in (827611204, 861165636, 894720068):   # DXT1 / DXT3 / DXT5
        bw = (w + 3) // 4
        bs = 8 if fmt == 827611204 else 16
        for by in range((h + 3) // 4):
            for bx in range(bw):
                boff = (by * bw + bx) * bs
                block = data[boff:boff+bs]
                if len(block) < bs:
                    continue
                coff = bs - 8
                c0, c1 = struct.unpack_from("<HH", block, coff)
                bits, = struct.unpack_from("<I", block, coff + 4)
                cols = []
                for c in (c0, c1):
                    cols.append((((c >> 11) & 31) * 255 // 31,
                                 ((c >> 5) & 63) * 255 // 63,
                                 (c & 31) * 255 // 31))
                if c0 > c1 or bs == 16:
                    cols.append(tuple((2*a+b)//3 for a, b in zip(cols[0], cols[1])))
                    cols.append(tuple((a+2*b)//3 for a, b in zip(cols[0], cols[1])))
                else:
                    cols.append(tuple((a+b)//2 for a, b in zip(cols[0], cols[1])))
                    cols.append((0, 0, 0))
                for py in range(4):
                    for px in range(4):
                        x, y = bx*4+px, by*4+py
                        if x >= w or y >= h:
                            continue
                        sel = (bits >> ((py*4+px)*2)) & 3
                        r, g, b = cols[sel]
                        if fmt == 827611204:
                            a = 0 if (sel == 3 and c0 <= c1) else 255
                        elif fmt == 861165636:      # DXT3 explicit alpha
                            av = block[(py*4+px)//2]
                            a = ((av >> 4) if (px & 1) else (av & 15)) * 17
                        else:                        # DXT5 interpolated alpha
                            a0, a1 = block[0], block[1]
                            abits = int.from_bytes(block[2:8], "little")
                            sel_a = (abits >> ((py*4+px)*3)) & 7
                            if sel_a == 0: a = a0
                            elif sel_a == 1: a = a1
                            elif a0 > a1:
                                a = ((8-sel_a)*a0 + (sel_a-1)*a1) // 7
                            elif sel_a == 6: a = 0
                            elif sel_a == 7: a = 255
                            else:
                                a = ((6-sel_a)*a0 + (sel_a-1)*a1) // 5
                        i = (y*w + x) * 4
                        out[i:i+4] = bytes((r, g, b, a))
    else:
        return None
    return bytes(out)


class Extract(unpickle.Unpickler):
    def __init__(self, stream):
        super().__init__(stream)
        self.blobs = Blobs()
        self.vb = {}          # ptr -> {"len":n,"data":bytearray,"dirty":True,"blob":id}
        self.ib = {}
        self.tex = {}         # ptr -> {"w","h","fmt","levels","level0":bytes,"pitch":..,"dirty"}
        self.surf2tex = {}
        self.locked = []      # (pbits, kind, owner, extra)
        self.state = {}
        self.tss = {}
        self.transforms = {}  # 2 view, 3 proj, 256 world
        self.fvf = 0
        self.stream0 = (0, 0, 0)
        self.indices = 0
        self.textures = {0: 0, 1: 0}
        self.frame_draws = 0
        self.recording = False
        self.frames = []
        self.cur = None
        self.tex_emitted = {}
        self.viewport = None
        self.clear_color = 0
        self.calls = 0
        self.material = None
        self.lights = {}
        self.lights_on = {}

    # -- helpers
    def argmap(self, call):
        return dict(call.args)

    def find_lock(self, dest):
        best = None
        for entry in self.locked:
            if entry[0] <= dest and (best is None or entry[0] > best[0]):
                best = entry
        return best

    def vb_slice(self, ptr, byte_off, byte_len):
        e = self.vb.get(ptr)
        if e is None or byte_len <= 0:
            return -1
        return self.blobs.put(bytes(e["data"][byte_off:byte_off+byte_len]))

    def ib_slice(self, ptr, first_index, index_count):
        e = self.ib.get(ptr)
        if e is None or index_count <= 0:
            return -1
        return self.blobs.put(bytes(
            e["data"][first_index*2:(first_index+index_count)*2]))

    def tex_ref(self, ptr):
        if not ptr or ptr not in self.tex:
            return -1
        t = self.tex[ptr]
        key = ptr
        if key not in self.tex_emitted or t["dirty"]:
            rgba = decode_texture(t["fmt"], t["w"], t["h"],
                                  t.get("level0", b""), t.get("pitch", t["w"]*4))
            if rgba is None:
                return -1
            self.tex_emitted[key] = {"w": t["w"], "h": t["h"],
                                     "blob": self.blobs.put(rgba)}
            t["dirty"] = False
        return list(self.tex_emitted).index(key)

    # -- main
    def handleCall(self, c):
        self.calls += 1
        n = c.functionName.split("::")[-1]
        iface = c.functionName.split("::")[0]
        a = self.argmap(c)

        if n == "CreateDevice" or n == "Reset":
            pp = a.get("pPresentationParameters")
            while isinstance(pp, list) and len(pp) == 1:
                pp = pp[0]
            if isinstance(pp, dict):
                self.bbw = pp.get("BackBufferWidth", 960)
                self.bbh = pp.get("BackBufferHeight", 720)
        elif n == "CreateVertexBuffer":
            self.vb[int(a["ppVertexBuffer"][0])] = {
                "len": a["Length"], "data": bytearray(a["Length"]),
                "dirty": True, "blob": -1}
        elif n == "CreateIndexBuffer":
            self.ib[int(a["ppIndexBuffer"][0])] = {
                "len": a["Length"], "data": bytearray(a["Length"]),
                "dirty": True, "blob": -1}
        elif n == "CreateTexture":
            self.tex[int(a["ppTexture"][0])] = {
                "w": a["Width"], "h": a["Height"], "fmt": a["Format"],
                "levels": a["Levels"], "dirty": True}
        elif n == "GetSurfaceLevel":
            self.surf2tex[int(a["ppSurfaceLevel"][0])] = (int(a["this"]), a["Level"])
        elif n == "Lock" and iface == "IDirect3DVertexBuffer9":
            self.locked.append((int(a["ppbData"][0]), "vb", int(a["this"]),
                                a.get("OffsetToLock", 0)))
        elif n == "Lock" and iface == "IDirect3DIndexBuffer9":
            self.locked.append((int(a["ppbData"][0]), "ib", int(a["this"]),
                                a.get("OffsetToLock", 0)))
        elif n == "LockRect":
            rect = a["pLockedRect"][0]
            owner = int(a["this"])
            level = a.get("Level", None)
            if iface == "IDirect3DSurface9":
                tex = self.surf2tex.get(owner)
                if tex:
                    owner, level = tex
                else:
                    return
            self.locked.append((int(rect["pBits"]), "tex", owner,
                                (level or 0, rect["Pitch"])))
        elif n in ("Unlock", "UnlockRect"):
            owner = int(a["this"])
            if iface == "IDirect3DSurface9":
                st = self.surf2tex.get(owner)
                owner = st[0] if st else owner
            self.locked = [l for l in self.locked if l[2] != owner]
        elif n == "memcpy":
            hit = self.find_lock(int(a["dest"]))
            if not hit:
                return
            base, kind, owner, extra = hit
            data = a["src"]
            if not isinstance(data, (bytes, bytearray)):
                return
            off = int(a["dest"]) - base
            if kind == "vb" and owner in self.vb:
                e = self.vb[owner]
                lockoff = extra + off
                e["data"][lockoff:lockoff+len(data)] = data
                e["dirty"] = True
            elif kind == "ib" and owner in self.ib:
                e = self.ib[owner]
                lockoff = extra + off
                e["data"][lockoff:lockoff+len(data)] = data
                e["dirty"] = True
            elif kind == "tex" and owner in self.tex:
                level, pitch = extra
                if level == 0 and off == 0:
                    self.tex[owner]["level0"] = bytes(data)
                    self.tex[owner]["pitch"] = pitch
                    self.tex[owner]["dirty"] = True
        elif n == "SetFVF":
            self.fvf = a["FVF"]
        elif n == "SetStreamSource":
            if a["StreamNumber"] == 0:
                self.stream0 = (int(a["pStreamData"]), a["OffsetInBytes"],
                                a["Stride"])
        elif n == "SetIndices":
            self.indices = int(a["pIndexData"])
        elif n == "SetTexture":
            self.textures[a["Stage"]] = int(a["pTexture"])
        elif n == "SetTransform":
            m = a["pMatrix"]
            while isinstance(m, list) and len(m) == 1:
                m = m[0]
            if isinstance(m, dict):
                m = m.get("m", m.get("_11", m))
            while isinstance(m, list) and len(m) == 1:
                m = m[0]
            flat = []
            if isinstance(m, list):
                for row in m:
                    flat.extend(row if isinstance(row, list) else [row])
            if len(flat) == 16:
                self.transforms[a["State"]] = [float(x) for x in flat]
        elif n == "SetMaterial":
            m = a.get("pMaterial")
            while isinstance(m, list) and len(m) == 1:
                m = m[0]
            if isinstance(m, dict):
                def col(c):
                    if isinstance(c, dict):
                        return [c.get("r", 1), c.get("g", 1), c.get("b", 1),
                                c.get("a", 1)]
                    return [1, 1, 1, 1]
                self.material = {"d": col(m.get("Diffuse")),
                                 "a": col(m.get("Ambient")),
                                 "e": col(m.get("Emissive"))}
        elif n == "SetLight":
            l = a.get("pLight")
            while isinstance(l, list) and len(l) == 1:
                l = l[0]
            if isinstance(l, dict):
                def vec(v):
                    if isinstance(v, dict):
                        return [v.get("x", 0), v.get("y", 0), v.get("z", 0)]
                    return [0, 0, -1]
                def col(c):
                    if isinstance(c, dict):
                        return [c.get("r", 1), c.get("g", 1), c.get("b", 1)]
                    return [1, 1, 1]
                self.lights[a.get("Index", 0)] = {
                    "type": l.get("Type", 3), "dir": vec(l.get("Direction")),
                    "dif": col(l.get("Diffuse")), "amb": col(l.get("Ambient"))}
        elif n == "LightEnable":
            self.lights_on[a.get("Index", 0)] = bool(a.get("Enable", 0))
        elif n == "SetRenderState":
            if a["State"] in REL_STATES:
                self.state[a["State"]] = a["Value"]
        elif n == "SetTextureStageState":
            if a["Stage"] <= 1:
                self.tss[f'{a["Stage"]}_{a["Type"]}'] = a["Value"]
        elif n == "Clear":
            self.clear_color = a.get("Color", 0)
        elif n == "SetViewport":
            self.viewport = a["pViewport"][0] if a.get("pViewport") else None
        elif n in DRAWS:
            self.frame_draws += 1
            if self.recording and self.cur is not None:
                stride = self.stream0[2]
                d = {"fn": n, "fvf": self.fvf, "stride": stride,
                     "st": dict(self.state),
                     "wm": self.transforms.get(256), "vm": self.transforms.get(2),
                     "pm": self.transforms.get(3),
                     "tm": self.transforms.get(16),
                     "mat": self.material,
                     "tex0": self.tex_ref(self.textures.get(0, 0))}
                if self.state.get(137):
                    active = [self.lights[i] for i, on in self.lights_on.items()
                              if on and i in self.lights
                              and self.lights[i]["type"] == 3]
                    if active:
                        d["light"] = active[0]
                if n == "DrawIndexedPrimitive" and stride:
                    prim, count = a["PrimitiveType"], a["primCount"]
                    nidx = count * 3 if prim == 4 else count + 2
                    vfirst = a["BaseVertexIndex"] + a["MinVertexIndex"]
                    vcount = a["NumVertices"]
                    d.update({"prim": prim, "count": count,
                              "minv": a["MinVertexIndex"],
                              "vb": self.vb_slice(self.stream0[0],
                                    self.stream0[1] + a["BaseVertexIndex"]*stride,
                                    (a["MinVertexIndex"]+vcount)*stride),
                              "ib": self.ib_slice(self.indices,
                                                  a["startIndex"], nidx)})
                elif n == "DrawPrimitive" and stride:
                    prim, count = a["PrimitiveType"], a["PrimitiveCount"]
                    nvert = count * 3 if prim == 4 else count + 2
                    d.update({"prim": prim, "count": count,
                              "vb": self.vb_slice(self.stream0[0],
                                    self.stream0[1] + a["StartVertex"]*stride,
                                    nvert*stride)})
                else:
                    return          # skip UP draws in v1
                self.cur["draws"].append(d)
        elif n == "Present":
            self.present_count = getattr(self, "present_count", 0) + 1
            if self.recording:
                if self.cur is not None:
                    self.frames.append(self.cur)
                    if len(self.frames) >= RECORD_FRAMES:
                        raise StopIteration
                self.cur = {"clear": self.clear_color, "draws": []}
            elif self.present_count >= START_FRAME:
                self.recording = True
                self.cur = {"clear": self.clear_color, "draws": []}
                print(f"recording from call {c.no} (frame {self.present_count})",
                      file=sys.stderr)
            self.frame_draws = 0


ex = Extract(sys.stdin.buffer)
try:
    ex.parse()
except StopIteration:
    pass

textures = [{"w": v["w"], "h": v["h"], "blob": v["blob"]}
            for v in ex.tex_emitted.values()]
bundle = {"width": getattr(ex, "bbw", 960), "height": getattr(ex, "bbh", 720),
          "textures": textures,
          "blobs": ex.blobs.entries, "frames": ex.frames}
with open(os.path.join(OUT, "bundle.json"), "w") as f:
    json.dump(bundle, f)
ex.blobs.f.close()

total_draws = sum(len(f["draws"]) for f in ex.frames)
print(f"frames={len(ex.frames)} draws={total_draws} textures={len(textures)}")
print(f"blob store={ex.blobs.off/1e6:.1f}MB unique / {ex.blobs.raw_bytes/1e6:.1f}MB raw "
      f"(dedup {100*(1-ex.blobs.off/max(1,ex.blobs.raw_bytes)):.1f}%)")
print(f"bundle.json={os.path.getsize(os.path.join(OUT,'bundle.json'))/1e6:.1f}MB")