#!/usr/bin/env python3
"""Convert a W3CS recorder capture into the WebGPU proof-player bundle."""

from __future__ import annotations

import argparse
from collections import Counter
import hashlib
import importlib.util
import json
from pathlib import Path
import struct
import sys


ROOT = Path(__file__).resolve().parents[1]
SPEC = importlib.util.spec_from_file_location("w3cs_protocol", ROOT / "protocol.py")
protocol = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = protocol
SPEC.loader.exec_module(protocol)

OP_FRAME_STATE = 1
OP_CREATE_BUFFER = 2
OP_CREATE_TEXTURE = 4
OP_UPDATE_TEXTURE = 5
OP_DESTROY_RESOURCE = 6
OP_DEFINE_BLOB = 7
OP_DEFINE_BLOB_DELTA = 8
OP_DEFINE_BLOB_XOR_MASK = 9
OP_RESET_BLOB_CACHE = 10
OP_SET_RENDER_STATE = 16
OP_SET_TRANSFORM = 17
OP_SET_WORLD_TRANSFORM_COMPACT = 34
OP_SET_TEXTURE = 18
OP_SET_FVF = 19
OP_SET_STREAM_SOURCE = 20
OP_SET_INDICES = 21
OP_SET_MATERIAL = 22
OP_SET_LIGHT = 23
OP_LIGHT_ENABLE = 24
OP_SET_TEXTURE_STAGE_STATE = 25
OP_SET_SAMPLER_STATE = 26
OP_SET_VIEWPORT = 27
OP_SET_SCISSOR = 28
OP_CLEAR = 29
OP_DRAW_PRIMITIVE = 32
OP_DRAW_INDEXED_PRIMITIVE = 33

DXT1 = 827611204
DXT3 = 861165636
DXT5 = 894720068
RELEVANT_RENDER_STATES = {
    7, 8, 9, 14, 15, 19, 20, 22, 23, 24, 25, 26, 27, 28, 29,
    34, 35, 36, 37, 137, 139, 140, 141, 145, 168,
}


def texture_layout(fmt: int, width: int, height: int) -> tuple[int, int]:
    if fmt == DXT1:
        return ((width + 3) // 4) * 8, (height + 3) // 4
    if fmt in (DXT3, DXT5):
        return ((width + 3) // 4) * 16, (height + 3) // 4
    bpp = 4 if fmt in (21, 22) else 2
    return width * bpp, height


def decode_texture(fmt: int, width: int, height: int, data: bytes) -> bytes | None:
    pitch, _ = texture_layout(fmt, width, height)
    out = bytearray(width * height * 4)
    if fmt in (21, 22):
        for y in range(height):
            for x in range(width):
                b, g, r, a = data[y * pitch + x * 4:y * pitch + x * 4 + 4]
                i = (y * width + x) * 4
                out[i:i + 4] = bytes((r, g, b, 255 if fmt == 22 else a))
    elif fmt in (23, 24, 25, 26):
        for y in range(height):
            for x in range(width):
                value, = struct.unpack_from("<H", data, y * pitch + x * 2)
                i = (y * width + x) * 4
                if fmt == 23:
                    rgba = (((value >> 11) & 31) * 255 // 31,
                            ((value >> 5) & 63) * 255 // 63,
                            (value & 31) * 255 // 31, 255)
                elif fmt in (24, 25):
                    alpha = 255 if fmt == 24 or value & 0x8000 else 0
                    rgba = (((value >> 10) & 31) * 255 // 31,
                            ((value >> 5) & 31) * 255 // 31,
                            (value & 31) * 255 // 31, alpha)
                else:
                    rgba = (((value >> 8) & 15) * 17,
                            ((value >> 4) & 15) * 17,
                            (value & 15) * 17, ((value >> 12) & 15) * 17)
                out[i:i + 4] = bytes(rgba)
    elif fmt in (DXT1, DXT3, DXT5):
        block_width = (width + 3) // 4
        block_size = 8 if fmt == DXT1 else 16
        for block_y in range((height + 3) // 4):
            for block_x in range(block_width):
                offset = (block_y * block_width + block_x) * block_size
                block = data[offset:offset + block_size]
                if len(block) != block_size:
                    continue
                color_offset = block_size - 8
                c0, c1, bits = struct.unpack_from("<HHI", block, color_offset)
                colors = []
                for color in (c0, c1):
                    colors.append((((color >> 11) & 31) * 255 // 31,
                                   ((color >> 5) & 63) * 255 // 63,
                                   (color & 31) * 255 // 31))
                if c0 > c1 or block_size == 16:
                    colors.append(tuple((2 * a + b) // 3
                                        for a, b in zip(colors[0], colors[1])))
                    colors.append(tuple((a + 2 * b) // 3
                                        for a, b in zip(colors[0], colors[1])))
                else:
                    colors.append(tuple((a + b) // 2
                                        for a, b in zip(colors[0], colors[1])))
                    colors.append((0, 0, 0))
                for py in range(4):
                    for px in range(4):
                        x, y = block_x * 4 + px, block_y * 4 + py
                        if x >= width or y >= height:
                            continue
                        selector = (bits >> ((py * 4 + px) * 2)) & 3
                        r, g, b = colors[selector]
                        if fmt == DXT1:
                            alpha = 0 if selector == 3 and c0 <= c1 else 255
                        elif fmt == DXT3:
                            alpha_bits = block[(py * 4 + px) // 2]
                            alpha = ((alpha_bits >> 4) if px & 1
                                     else (alpha_bits & 15)) * 17
                        else:
                            a0, a1 = block[0], block[1]
                            alpha_bits = int.from_bytes(block[2:8], "little")
                            alpha_selector = (alpha_bits >> ((py * 4 + px) * 3)) & 7
                            if alpha_selector == 0:
                                alpha = a0
                            elif alpha_selector == 1:
                                alpha = a1
                            elif a0 > a1:
                                alpha = ((8 - alpha_selector) * a0
                                         + (alpha_selector - 1) * a1) // 7
                            elif alpha_selector == 6:
                                alpha = 0
                            elif alpha_selector == 7:
                                alpha = 255
                            else:
                                alpha = ((6 - alpha_selector) * a0
                                         + (alpha_selector - 1) * a1) // 5
                        i = (y * width + x) * 4
                        out[i:i + 4] = bytes((r, g, b, alpha))
    else:
        return None
    return bytes(out)


class BlobWriter:
    def __init__(self, output: Path):
        self.file = (output / "blobs.bin").open("wb")
        self.index: dict[bytes, int] = {}
        self.entries: list[tuple[int, int]] = []
        self.offset = 0

    def put(self, data: bytes) -> int:
        digest = hashlib.blake2b(data, digest_size=16).digest()
        existing = self.index.get(digest)
        if existing is not None:
            return existing
        index = len(self.entries)
        self.index[digest] = index
        self.entries.append((self.offset, len(data)))
        self.file.write(data)
        self.offset += len(data)
        return index


class Converter:
    def __init__(self, output: Path, min_draws: int, frame_limit: int,
                 source_fps: int):
        self.output = output
        self.min_draws = min_draws
        self.frame_limit = frame_limit
        self.source_fps = source_fps
        self.blob_writer = BlobWriter(output)
        self.buffers: dict[int, dict] = {}
        self.textures: dict[int, dict] = {}
        self.command_blobs: dict[int, bytes] = {}
        self.command_blob_bases: dict[int, int] = {}
        self.texture_versions: dict[tuple[int, int], int] = {}
        self.output_textures: list[dict] = []
        self.frames: list[dict] = []
        self.texture_draws = Counter()
        self.missing_texture_draws = Counter()
        self.reset_state()

    def reset_state(self):
        clear_color = getattr(self, "clear_color", 0)
        self.render_state = {}
        self.transforms = {}
        self.tss = {}
        self.samplers = {}
        self.texture_bindings = {}
        self.streams = {}
        self.index_binding = (0, 0)
        self.fvf = 0
        self.material = None
        self.lights = {}
        self.lights_enabled = {}
        self.viewport = None
        self.scissor = None
        self.clear_color = clear_color

    def apply_texture_update(self, payload: bytes):
        if len(payload) < 36:
            raise protocol.ProtocolError("short texture update")
        (resource_id, generation, level, x, y, width, height, pitch,
         fmt) = struct.unpack_from("<9I", payload)
        size, = struct.unpack_from("<I", payload, 36)
        data = payload[40:]
        if len(data) != size or level != 0:
            return
        texture = self.textures.get(resource_id)
        if not texture or texture["fmt"] != fmt:
            return
        dest_pitch, dest_rows = texture_layout(fmt, texture["w"], texture["h"])
        if "data" not in texture:
            texture["data"] = bytearray(dest_pitch * dest_rows)
        if fmt in (DXT1, DXT3, DXT5):
            block_size = 8 if fmt == DXT1 else 16
            dest_y, dest_x = y // 4, (x // 4) * block_size
            rows = (height + 3) // 4
        else:
            bpp = 4 if fmt in (21, 22) else 2
            dest_y, dest_x, rows = y, x * bpp, height
        for row in range(rows):
            source_start = row * pitch
            dest_start = (dest_y + row) * dest_pitch + dest_x
            texture["data"][dest_start:dest_start + pitch] = data[
                source_start:source_start + pitch]
        texture["generation"] = generation
        for stage, binding in list(self.texture_bindings.items()):
            if binding[0] == resource_id:
                self.texture_bindings[stage] = (resource_id, generation)

    def texture_ref(self, resource_id: int, generation: int) -> int:
        if not resource_id:
            return -1
        key = (resource_id, generation)
        existing = self.texture_versions.get(key)
        if existing is not None:
            return existing
        texture = self.textures.get(resource_id)
        if not texture or "data" not in texture:
            return -1
        rgba = decode_texture(texture["fmt"], texture["w"], texture["h"],
                              bytes(texture["data"]))
        if rgba is None:
            return -1
        index = len(self.output_textures)
        self.output_textures.append({
            "w": texture["w"], "h": texture["h"],
            "blob": self.blob_writer.put(rgba),
        })
        self.texture_versions[key] = index
        return index

    def draw_common(self, vertex_blob: int) -> dict:
        stream = self.streams.get(0, (0, 0, 0, 0))
        texture0 = self.texture_bindings.get(0, (0, 0))
        texture1 = self.texture_bindings.get(1, (0, 0))
        self.texture_draws[texture0] += 1
        texture0_index = self.texture_ref(*texture0)
        texture1_index = self.texture_ref(*texture1)
        if texture0_index < 0 and texture0[0]:
            self.missing_texture_draws[texture0] += 1
        base_id = self.command_blob_bases.get(vertex_blob)
        draw = {
            "fvf": self.fvf,
            "stride": stream[3],
            "st": dict(self.render_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": texture0_index,
            "tex1": texture1_index,
            "tss": {f"{stage}_{state}": value
                    for (stage, state), value in self.tss.items()
                    if stage <= 1},
            "samp": {f"{stage}_{state}": value
                     for (stage, state), value in self.samplers.items()
                     if stage <= 1},
            "vb": self.blob_writer.put(self.command_blobs[vertex_blob])
                  if vertex_blob in self.command_blobs else -1,
            "vbbase": self.blob_writer.put(self.command_blobs[base_id])
                      if base_id in self.command_blobs else -1,
        }
        if self.render_state.get(137):
            active = [self.lights[index] for index, enabled
                      in self.lights_enabled.items()
                      if enabled and index in self.lights
                      and self.lights[index]["type"] == 3]
            if active:
                draw["light"] = active[0]
        return draw

    @staticmethod
    def parse_material(payload: bytes) -> dict | None:
        if len(payload) != 68:
            return None
        values = struct.unpack("<17f", payload)
        return {"d": list(values[0:4]), "a": list(values[4:8]),
                "e": list(values[12:16])}

    @staticmethod
    def parse_light(payload: bytes) -> tuple[int, dict] | None:
        if len(payload) != 108:
            return None
        index, light_type = struct.unpack_from("<II", payload)
        values = struct.unpack_from("<25f", payload, 8)
        return index, {"type": light_type, "dif": list(values[0:3]),
                       "amb": list(values[8:11]), "dir": list(values[15:18])}

    def apply_records(self, records: list[protocol.Record], is_frame: bool):
        frame = {"clear": self.clear_color, "draws": []} if is_frame else None
        for record in records:
            payload = record.payload
            opcode = record.opcode
            if opcode == OP_FRAME_STATE:
                self.reset_state()
                if frame is not None:
                    frame["clear"] = self.clear_color
            elif opcode == OP_CREATE_BUFFER and len(payload) == 25:
                rid, generation, size, usage, fmt, pool = struct.unpack_from(
                    "<6I", payload)
                self.buffers[rid] = {"generation": generation, "size": size,
                                     "usage": usage, "format": fmt,
                                     "pool": pool, "kind": payload[24]}
            elif opcode == OP_CREATE_TEXTURE and len(payload) == 32:
                rid, generation, width, height, levels, usage, fmt, pool = \
                    struct.unpack("<8I", payload)
                self.textures[rid] = {"generation": generation, "w": width,
                                      "h": height, "levels": levels,
                                      "usage": usage, "fmt": fmt, "pool": pool}
            elif opcode == OP_UPDATE_TEXTURE:
                self.apply_texture_update(payload)
            elif opcode == OP_DESTROY_RESOURCE and len(payload) == 4:
                rid, = struct.unpack("<I", payload)
                self.buffers.pop(rid, None)
                self.textures.pop(rid, None)
            elif opcode == OP_DEFINE_BLOB and len(payload) >= 24:
                rid, size = struct.unpack_from("<II", payload)
                if size == len(payload) - 24:
                    self.command_blobs[rid] = bytes(payload[24:])
                    self.command_blob_bases.pop(rid, None)
            elif opcode == OP_DEFINE_BLOB_DELTA and len(payload) >= 28:
                rid, size = struct.unpack_from("<II", payload)
                base_id, = struct.unpack_from("<I", payload, 24)
                base = self.command_blobs.get(base_id)
                delta = payload[28:]
                if base is not None and len(base) == size and len(delta) == size:
                    self.command_blobs[rid] = bytes(
                        left ^ right for left, right in zip(base, delta))
                    self.command_blob_bases[rid] = base_id
            elif opcode == OP_DEFINE_BLOB_XOR_MASK and len(payload) >= 28:
                rid, size = struct.unpack_from("<II", payload)
                base_id, = struct.unpack_from("<I", payload, 24)
                base = self.command_blobs.get(base_id)
                word_count = (size + 3) // 4
                mask_size = (word_count + 1) // 2
                mask = payload[28:28 + mask_size]
                values = payload[28 + mask_size:]
                if base is not None and len(base) == size \
                        and len(mask) == mask_size:
                    output = bytearray(base)
                    value_offset = 0
                    valid = True
                    for index in range(size):
                        word = index // 4
                        bits = mask[word // 2] >> (4 if word & 1 else 0)
                        if bits & (1 << (index & 3)):
                            if value_offset >= len(values):
                                valid = False
                                break
                            output[index] ^= values[value_offset]
                            value_offset += 1
                    if valid and value_offset == len(values):
                        self.command_blobs[rid] = bytes(output)
                        self.command_blob_bases[rid] = base_id
            elif opcode == OP_RESET_BLOB_CACHE and not payload:
                self.command_blobs.clear()
                self.command_blob_bases.clear()
            elif opcode == OP_SET_RENDER_STATE and len(payload) == 8:
                state, value = struct.unpack("<II", payload)
                if state in RELEVANT_RENDER_STATES:
                    self.render_state[state] = value
            elif opcode == OP_SET_TRANSFORM and len(payload) == 68:
                state, = struct.unpack_from("<I", payload)
                self.transforms[state] = list(struct.unpack_from("<16f", payload, 4))
            elif opcode == OP_SET_WORLD_TRANSFORM_COMPACT and len(payload) == 38:
                halves = struct.unpack_from("<13e", payload)
                matrix = list(halves[:12]) + list(struct.unpack_from("<3f", payload, 26))
                matrix.append(halves[12])
                self.transforms[256] = matrix
            elif opcode == OP_SET_TEXTURE and len(payload) == 12:
                stage, rid, generation = struct.unpack("<III", payload)
                self.texture_bindings[stage] = (rid, generation)
            elif opcode == OP_SET_FVF and len(payload) == 4:
                self.fvf, = struct.unpack("<I", payload)
            elif opcode == OP_SET_STREAM_SOURCE and len(payload) == 20:
                stage, rid, generation, offset, stride = struct.unpack("<5I", payload)
                self.streams[stage] = (rid, generation, offset, stride)
            elif opcode == OP_SET_INDICES and len(payload) == 8:
                self.index_binding = struct.unpack("<II", payload)
            elif opcode == OP_SET_MATERIAL:
                self.material = self.parse_material(payload)
            elif opcode == OP_SET_LIGHT:
                light = self.parse_light(payload)
                if light:
                    self.lights[light[0]] = light[1]
            elif opcode == OP_LIGHT_ENABLE and len(payload) == 8:
                index, enabled = struct.unpack("<II", payload)
                self.lights_enabled[index] = bool(enabled)
            elif opcode == OP_SET_TEXTURE_STAGE_STATE and len(payload) == 12:
                stage, state, value = struct.unpack("<III", payload)
                self.tss[(stage, state)] = value
            elif opcode == OP_SET_SAMPLER_STATE and len(payload) == 12:
                stage, state, value = struct.unpack("<III", payload)
                self.samplers[(stage, state)] = value
            elif opcode == OP_SET_VIEWPORT and len(payload) == 24:
                self.viewport = struct.unpack("<4I2f", payload)
            elif opcode == OP_SET_SCISSOR and len(payload) == 16:
                self.scissor = struct.unpack("<4i", payload)
            elif opcode == OP_CLEAR and len(payload) == 20:
                _, _, self.clear_color = struct.unpack_from("<3I", payload)
                if frame is not None:
                    frame["clear"] = self.clear_color
            elif opcode == OP_DRAW_PRIMITIVE and frame is not None \
                    and len(payload) == 16:
                primitive, start, count, vertex_blob = struct.unpack("<4I", payload)
                draw = self.draw_common(vertex_blob)
                draw.update({"fn": "DrawPrimitive", "prim": primitive,
                             "count": count, "start": start})
                frame["draws"].append(draw)
            elif opcode == OP_DRAW_INDEXED_PRIMITIVE and frame is not None \
                    and len(payload) == 32:
                (primitive, base, minimum, vertices, start, count,
                 vertex_blob, index_blob) = struct.unpack("<Ii6I", payload)
                draw = self.draw_common(vertex_blob)
                draw.update({"fn": "DrawIndexedPrimitive", "prim": primitive,
                             "count": count, "base": base, "minv": minimum,
                             "numv": vertices, "start": start,
                             "ib": self.blob_writer.put(self.command_blobs[index_blob])
                                   if index_blob in self.command_blobs else -1,
                             "index32": self.buffers.get(
                                 self.index_binding[0], {}).get("format") == 102})
                frame["draws"].append(draw)
        if frame is not None and len(frame["draws"]) >= self.min_draws \
                and len(self.frames) < self.frame_limit:
            self.frames.append(frame)

    def convert(self, capture: Path):
        reassembler = protocol.Reassembler()
        with capture.open("rb") as source:
            while len(self.frames) < self.frame_limit:
                raw_size = source.read(4)
                if not raw_size:
                    break
                if len(raw_size) != 4:
                    raise protocol.ProtocolError("truncated packet length")
                size = int.from_bytes(raw_size, "little")
                packet = source.read(size)
                if len(packet) != size:
                    raise protocol.ProtocolError("truncated packet")
                completed = reassembler.push(packet)
                if completed is None:
                    continue
                envelope, payload = completed
                if envelope.kind in (protocol.Kind.RESOURCE, protocol.Kind.FRAME):
                    self.apply_records(protocol.decode_records(payload),
                                       envelope.kind is protocol.Kind.FRAME)

    def finish(self):
        self.blob_writer.file.close()
        bundle = {"width": 1024, "height": 768, "fps": self.source_fps,
                  "textures": self.output_textures,
                  "blobs": self.blob_writer.entries, "frames": self.frames}
        (self.output / "bundle.json").write_text(
            json.dumps(bundle, separators=(",", ":")))
        print(f"frames={len(self.frames)} "
              f"draws={sum(len(frame['draws']) for frame in self.frames)} "
              f"textures={len(self.output_textures)} "
              f"blobs={len(self.blob_writer.entries)} "
              f"bytes={self.blob_writer.offset}")
        if self.missing_texture_draws:
            missing = [(key, count, self.textures.get(key[0]))
                       for key, count in self.missing_texture_draws.most_common(12)]
            print("missing texture draws:", missing)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("capture", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--min-draws", type=int, default=100)
    parser.add_argument("--frames", type=int, default=240)
    parser.add_argument("--fps", type=int, default=40)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    converter = Converter(args.output, args.min_draws, args.frames, args.fps)
    converter.convert(args.capture)
    converter.finish()


if __name__ == "__main__":
    main()