#!/usr/bin/env python3
"""Compression closes the gap: raw vs gzipped encodings, DDIA 2e Ch. 5.

Binary formats (Protobuf/Avro) look 2-3x smaller than JSON when you measure the
RAW bytes. But most of JSON's bulk is repeated field names and punctuation --
exactly the redundancy gzip is built to remove. So we encode a batch of 10,000
records four ways, measure raw bytes, gzip each blob at level 9, and watch the
binary advantage largely COLLAPSE. The lesson: once you can compress a batch,
"smaller on the wire" is a weak reason to pick binary; schema/speed/typing are
the real reasons.

Run it:
    uv run --with msgpack --with fastavro --with grpcio-tools --with protobuf \
        python3 compression_gap.py
(Protobuf's .proto is compiled at runtime via grpc_tools, so no separate protoc
step is needed. The batch is deterministic -> the byte counts reproduce exactly.)
"""
import gzip
import io
import json
import os
import sys
import tempfile

N = 10_000

# A batch of deterministic-but-varied user records. Same seed data every run,
# so every byte count below reproduces exactly.
INTERESTS_POOL = ["daydreaming", "hacking", "climbing", "cooking",
                  "reading", "cycling", "photography", "music"]


def make_records():
    records = []
    for i in range(N):
        interests = [INTERESTS_POOL[i % len(INTERESTS_POOL)],
                     INTERESTS_POOL[(i * 3 + 1) % len(INTERESTS_POOL)]]
        records.append({
            "id": i,
            "userName": f"user{i}",
            "favoriteNumber": (i * 2654435761) % 100000,
            "interests": interests,
        })
    return records


def json_bytes(records):
    # Newline-delimited JSON: one compact JSON object per line, as you would
    # write to a log or a Kafka topic.
    out = io.BytesIO()
    for r in records:
        out.write(json.dumps(r, separators=(",", ":")).encode())
        out.write(b"\n")
    return out.getvalue()


def msgpack_bytes(records):
    import msgpack
    out = io.BytesIO()
    for r in records:
        out.write(msgpack.packb(r))
    return out.getvalue()


def protobuf_bytes(records):
    work = tempfile.mkdtemp()
    proto = os.path.join(work, "person.proto")
    with open(proto, "w") as f:
        f.write('syntax = "proto3";\n'
                'message Person {\n'
                '  int64  id = 1;\n'
                '  string user_name = 2;\n'
                '  int64  favorite_number = 3;\n'
                '  repeated string interests = 4;\n'
                '}\n')
    from grpc_tools import protoc
    protoc.main(["protoc", f"-I{work}", f"--python_out={work}", proto])
    sys.path.insert(0, work)
    import person_pb2
    from google.protobuf.internal.encoder import _EncodeVarint
    out = io.BytesIO()
    for r in records:
        p = person_pb2.Person(id=r["id"], user_name=r["userName"],
                              favorite_number=r["favoriteNumber"],
                              interests=r["interests"])
        data = p.SerializeToString()
        # length-delimited framing, as a stream of messages needs
        _EncodeVarint(out.write, len(data))
        out.write(data)
    return out.getvalue()


def avro_bytes(records):
    import fastavro
    schema = fastavro.parse_schema({
        "type": "record", "name": "Person", "fields": [
            {"name": "id", "type": "long"},
            {"name": "userName", "type": "string"},
            {"name": "favoriteNumber", "type": "long"},
            {"name": "interests", "type": {"type": "array", "items": "string"}},
        ]})
    out = io.BytesIO()
    for r in records:
        fastavro.schemaless_writer(out, schema, r)
    return out.getvalue()


def main():
    records = make_records()

    blobs = [
        ("JSON (ndjson)", json_bytes(records)),
        ("MessagePack", msgpack_bytes(records)),
        ("Protocol Buffers", protobuf_bytes(records)),
        ("Avro", avro_bytes(records)),
    ]

    json_raw = blobs[0][1]
    json_raw_len = len(json_raw)
    json_gz_len = len(gzip.compress(json_raw, 9))

    print(f"batch of {N:,} user records "
          "{id, userName, favoriteNumber, interests[2]}\n")
    print(f"{'format':18s} {'raw':>10s} {'vs JSON':>8s}   "
          f"{'gzip -9':>10s} {'vs JSON':>8s}")
    print("-" * 62)
    rows = []
    for name, blob in blobs:
        raw = len(blob)
        gz = len(gzip.compress(blob, 9))
        raw_ratio = raw / json_raw_len
        gz_ratio = gz / json_gz_len
        rows.append((name, raw, gz, raw_ratio, gz_ratio))
        print(f"{name:18s} {raw:>10,d} {raw_ratio:>7.2f}x   "
              f"{gz:>10,d} {gz_ratio:>7.2f}x")

    # The surprise, stated in one line: how big was the binary win before gzip,
    # and how small is it after?
    av_raw_ratio = rows[3][3]
    av_gz_ratio = rows[3][4]
    print(f"\nraw:  Avro is {1 / av_raw_ratio:.2f}x smaller than JSON "
          f"({av_raw_ratio:.2f}x its size).")
    print(f"gzip: Avro is {1 / av_gz_ratio:.2f}x smaller than JSON "
          f"({av_gz_ratio:.2f}x its size) -- the gap largely closed.")
    print(f"\ngzip shrinks JSON by {100 * (1 - json_gz_len / json_raw_len):.0f}% "
          f"({json_raw_len:,} -> {json_gz_len:,} bytes): most of JSON's bulk was "
          "repeated field names + punctuation, which is exactly what gzip removes.")


if __name__ == "__main__":
    main()