#!/usr/bin/env python3
"""Summarize fixed-function state used by draws in a W3CS capture."""

from __future__ import annotations

import argparse
from collections import Counter
from pathlib import Path
import struct

import protocol


def summarize(path: Path) -> None:
    reassembler = protocol.Reassembler()
    textures: dict[int, tuple[int, int, int]] = {}
    bindings: dict[int, int] = {}
    stages: dict[tuple[int, int], int] = {}
    samplers: dict[tuple[int, int], int] = {}
    render: dict[int, int] = {}
    fvf = 0
    draws = 0
    frames = 0
    primitive_counts: Counter[int] = Counter()
    format_counts: Counter[tuple[int, int, int]] = Counter()
    stage_counts: Counter[tuple[int, ...]] = Counter()
    sampler_counts: Counter[tuple[int, ...]] = Counter()
    render_counts: Counter[tuple[int, ...]] = Counter()
    render_values: dict[int, Counter[int]] = {}
    second_stage_draws = 0

    with path.open("rb") as source:
        while size_bytes := source.read(4):
            if len(size_bytes) != 4:
                break
            size = int.from_bytes(size_bytes, "little")
            packet = source.read(size)
            if len(packet) != size:
                break
            complete = reassembler.push(packet)
            if complete is None:
                continue
            envelope, payload = complete
            if envelope.kind not in (protocol.Kind.RESOURCE,
                                     protocol.Kind.FRAME):
                continue
            if envelope.kind is protocol.Kind.FRAME:
                frames += 1
            for record in protocol.decode_records(payload):
                data = record.payload
                if record.opcode == 1:
                    bindings.clear()
                    stages.clear()
                    samplers.clear()
                    render.clear()
                    fvf = 0
                elif record.opcode == 4 and len(data) == 32:
                    rid, _, width, height, _, _, fmt, _ = struct.unpack(
                        "<8I", data)
                    textures[rid] = (fmt, width, height)
                elif record.opcode == 16 and len(data) == 8:
                    state, value = struct.unpack("<2I", data)
                    render[state] = value
                elif record.opcode == 18 and len(data) == 12:
                    stage, rid, _ = struct.unpack("<3I", data)
                    bindings[stage] = rid
                elif record.opcode == 19 and len(data) == 4:
                    fvf, = struct.unpack("<I", data)
                elif record.opcode == 25 and len(data) == 12:
                    stage, state, value = struct.unpack("<3I", data)
                    stages[(stage, state)] = value
                elif record.opcode == 26 and len(data) == 12:
                    stage, state, value = struct.unpack("<3I", data)
                    samplers[(stage, state)] = value
                elif record.opcode in (32, 33) and len(data) >= 4:
                    draws += 1
                    primitive, = struct.unpack_from("<I", data)
                    primitive_counts[primitive] += 1
                    texture0 = bindings.get(0, 0)
                    texture1 = bindings.get(1, 0)
                    format_counts[textures.get(texture0, (0, 0, 0))] += 1
                    if texture1:
                        second_stage_draws += 1
                    stage_counts[tuple(stages.get((stage, state), 0)
                        for stage in (0, 1) for state in range(1, 7))] += 1
                    sampler_counts[tuple(samplers.get((stage, state), 0)
                        for stage in (0, 1) for state in range(1, 11))] += 1
                    render_counts[(fvf, render.get(7, 0), render.get(14, 0),
                        render.get(15, 0), render.get(19, 0),
                        render.get(20, 0), render.get(22, 0),
                        render.get(23, 0), render.get(27, 0))] += 1
                    for state, value in render.items():
                        render_values.setdefault(state, Counter())[value] += 1

    print(f"frames={frames} draws={draws} stage1_texture_draws={second_stage_draws}")
    print("primitives", primitive_counts.most_common())
    print("texture_formats", format_counts.most_common(16))
    print("texture_stages", stage_counts.most_common(16))
    print("samplers", sampler_counts.most_common(16))
    print("render_states", render_counts.most_common(16))
    print("render_values", {state: values.most_common()
          for state, values in sorted(render_values.items())})


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("capture", type=Path)
    summarize(parser.parse_args().capture)


if __name__ == "__main__":
    main()