import os
import random
import shutil
import sys
import tempfile
from compression import zstd  # Python 3.14+ stdlib

import lance
import pyarrow as pa

print(f"python {sys.version.split()[0]}  |  pylance {lance.__version__}")

rng = random.Random(0)
vocab = [f"token{i:04d}" for i in range(600)]
docs = [" ".join(rng.choice(vocab) for _ in range(7000)) for _ in range(300)]  # ~55 KiB each
arr = pa.array(docs, type=pa.large_string())


def written_bytes(level):
    md = {"lance-encoding:compression": "zstd", "lance-encoding:compression-level": str(level)}
    field = pa.field("v", pa.large_string(), metadata=md)
    tmp = tempfile.mkdtemp()
    p = os.path.join(tmp, "ds")
    lance.write_dataset(pa.table([arr], schema=pa.schema([field])), p, data_storage_version="2.1")
    n = sum(os.path.getsize(os.path.join(d, f)) for d, _, fs in os.walk(p) for f in fs)
    shutil.rmtree(tmp, ignore_errors=True)
    return n


l1, l22 = written_bytes(1), written_bytes(22)

blob = ("".join(docs)).encode()
c1 = len(zstd.compress(blob, level=1))
c22 = len(zstd.compress(blob, level=22))

lance_gain = (1 - l22 / l1) * 100
control_gain = (1 - c22 / c1) * 100
print(f"value size ~{len(docs[0])} bytes (> 32 KiB)")
print(f"Lance  zstd L1 : {l1:>12,} bytes")
print(f"Lance  zstd L22: {l22:>12,} bytes   (L1->L22 {lance_gain:+.1f}%)")
print(f"stdlib zstd control L1->L22: {c1:,} -> {c22:,}  ({control_gain:.0f}% smaller)")