#!/usr/bin/env python3
"""Varint & zigzag: why -1 costs 10 bytes as int64, but 1 as sint64. DDIA 2e Ch. 5.

Protobuf encodes integers as variable-length varints: a small positive int is
one byte. But a plain `int64`/`int32` encodes a *negative* value by two's-
complement sign extension to the full 64-bit width, so every negative number --
even -1 -- fills all ten varint bytes. The `sint64`/`sint32` types apply ZIGZAG
encoding, which maps small-magnitude negatives to small varints, so -1 is one
byte again. This script isolates each encoding in its own single-field message
and measures the real serialized length.

Run it:
    uv run --python 3.12 --with grpcio-tools --with protobuf \
        python3 varint_zigzag.py
(The .proto is compiled at runtime via grpc_tools, so no separate protoc step is
needed. Deterministic -> the byte counts reproduce exactly.)
"""
import os
import sys
import tempfile


def load_messages():
    """Compile a .proto at runtime and return its message classes."""
    work = tempfile.mkdtemp()
    proto = os.path.join(work, "nums.proto")
    with open(proto, "w") as f:
        f.write('syntax = "proto3";\n'
                'message Plain  { int64  a = 1; }\n'          # sign-extended
                'message Zig    { sint64 a = 1; }\n'          # zigzag
                'message PlainN { repeated int64  a = 1; }\n' # packed batch
                'message ZigN   { repeated sint64 a = 1; }\n')
    from grpc_tools import protoc
    protoc.main(["protoc", f"-I{work}", f"--python_out={work}", proto])
    sys.path.insert(0, work)
    import nums_pb2
    return nums_pb2


def main():
    pb = load_messages()

    # --- per-value byte counts: one number, one single-field message ---
    values = [1, 1337, -1, -1000, 2147483647, -2147483648]
    print("per-value encoded size (single-field message, so total = tag + payload):")
    print(f"    {'value':>12}   {'int64':>7}   {'sint64':>7}   {'ratio':>6}")
    for v in values:
        p = len(pb.Plain(a=v).SerializeToString())
        z = len(pb.Zig(a=v).SerializeToString())
        ratio = f"{p / z:.1f}x" if z else "-"
        print(f"    {v:>12}   {p:>5}b   {z:>5}b   {ratio:>6}")

    # -1 is the headline: int64 sign-extends to 10 payload bytes (11 total).
    # Every single-field message spends 1 byte on the field tag (0x08 here),
    # so the varint PAYLOAD is total - 1: 10 bytes for int64, 1 byte for sint64.
    neg1_plain = len(pb.Plain(a=-1).SerializeToString())
    neg1_zig = len(pb.Zig(a=-1).SerializeToString())
    print(f"\nthe headline (-1): int64 = 1 tag + {neg1_plain - 1} varint payload "
          f"= {neg1_plain} bytes;  sint64 = 1 tag + {neg1_zig - 1} varint payload "
          f"= {neg1_zig} bytes.")
    print(f"    the varint payload alone: {neg1_plain - 1} bytes vs "
          f"{neg1_zig - 1} byte -- a {(neg1_plain - 1) // (neg1_zig - 1)}x blowup "
          f"for one sign bit.")

    # --- realistic batch: 1000 small negatives, packed repeated field ---
    batch = [-(i % 100) - 1 for i in range(1000)]  # small negatives, -1..-100
    big_plain = len(pb.PlainN(a=batch).SerializeToString())
    big_zig = len(pb.ZigN(a=batch).SerializeToString())
    print(f"\nbatch of {len(batch)} small negative ints (packed repeated field):")
    print(f"    repeated int64    {big_plain:>6} bytes")
    print(f"    repeated sint64   {big_zig:>6} bytes")
    print(f"    -> the plain int64 batch is {big_plain / big_zig:.1f}x larger, "
          f"wasting {big_plain - big_zig} bytes on sign extension.")


if __name__ == "__main__":
    main()