#!/usr/bin/env python3
"""Measure independent and streaming compression on a recorder capture."""

from __future__ import annotations

import argparse
import ctypes
import ctypes.util
import importlib.util
from pathlib import Path
import struct
import sys
import zlib


ROOT = Path(__file__).resolve().parent
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)


class Zstd:
    """Small dependency-free binding for one-shot libzstd measurements."""

    def __init__(self):
        name = ctypes.util.find_library("zstd")
        if not name:
            raise RuntimeError("libzstd is not installed")
        self.lib = ctypes.CDLL(name)
        self.lib.ZSTD_compressBound.argtypes = [ctypes.c_size_t]
        self.lib.ZSTD_compressBound.restype = ctypes.c_size_t
        self.lib.ZSTD_compress.argtypes = [
            ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p,
            ctypes.c_size_t, ctypes.c_int]
        self.lib.ZSTD_compress.restype = ctypes.c_size_t
        self.lib.ZSTD_isError.argtypes = [ctypes.c_size_t]
        self.lib.ZSTD_isError.restype = ctypes.c_uint

    def size(self, payload: bytes, level: int = 1) -> int:
        source = bytes(payload)
        bound = self.lib.ZSTD_compressBound(len(source))
        destination = ctypes.create_string_buffer(bound)
        result = self.lib.ZSTD_compress(
            destination, bound, source, len(source), level)
        if self.lib.ZSTD_isError(result):
            raise RuntimeError("libzstd failed to compress a payload")
        return int(result)


def measure(path: Path):
    zstd = Zstd()
    reassembler = protocol.Reassembler()
    independent = 0
    resource_stream_bytes = 0
    frame_bytes = 0
    resource_raw = 0
    frame_raw = 0
    compressor = zlib.compressobj(level=1)
    zstd_independent = 0
    zstd_batched = 0
    resource_batch = bytearray()

    def flush_resource_batch():
        nonlocal zstd_batched
        if resource_batch:
            zstd_batched += zstd.size(resource_batch)
            resource_batch.clear()

    with path.open("rb") as source:
        while raw_size := source.read(4):
            if len(raw_size) != 4:
                raise protocol.ProtocolError("truncated length")
            size, = struct.unpack("<I", raw_size)
            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 is protocol.Kind.RESOURCE:
                resource_raw += len(payload)
                independent += len(zlib.compress(payload, level=1))
                zstd_independent += zstd.size(payload)
                framed = struct.pack("<I", len(payload)) + payload
                resource_stream_bytes += len(compressor.compress(framed))
                resource_batch.extend(payload)
            elif envelope.kind is protocol.Kind.FRAME:
                flush_resource_batch()
                frame_raw += len(payload)
                encoded = zlib.compress(payload, level=1)
                independent += len(encoded)
                frame_bytes += len(encoded)
                frame_zstd = zstd.size(payload)
                zstd_independent += frame_zstd
                zstd_batched += frame_zstd
    flush_resource_batch()
    resource_stream_bytes += len(compressor.flush())
    print(f"resource_raw={resource_raw} frame_raw={frame_raw}")
    print(f"independent_deflate1={independent}")
    print("streaming_resources_plus_independent_frames="
          f"{resource_stream_bytes + frame_bytes} "
          f"resources={resource_stream_bytes} frames={frame_bytes}")
    print(f"independent_zstd1={zstd_independent}")
    print(f"batched_resources_plus_independent_frames_zstd1={zstd_batched}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("capture", type=Path)
    args = parser.parse_args()
    measure(args.capture)


if __name__ == "__main__":
    main()