#!/usr/bin/env python3
"""Per-tile zstd CLI comparison over the pristine L1 records from this round."""

from pathlib import Path
import shutil
import struct
import subprocess
import time
import zlib


REPO = Path.cwd().resolve()
CACHE = REPO / "Saved/Linux/cache/l16-prototype"
CAPTURE_LOG = REPO / "Saved/Linux/l16-cache-l1-q24-full-write.log"
TEMP = CACHE / "codec-zstd"
RESULT = REPO / "Saved/Linux/l16-zstd-codec.txt"

assert CACHE.is_dir() and not CACHE.is_symlink()
assert CACHE.resolve().parent == (REPO / "Saved/Linux/cache").resolve()
assert CAPTURE_LOG.is_file()
assert not TEMP.exists()

cutoff_ns = CAPTURE_LOG.stat().st_mtime_ns
records = [
    path for world in CACHE.glob("world_*")
    for path in world.glob("*.vftile")
    if "_s0_d0000000000000000_" in path.name and path.stat().st_mtime_ns <= cutoff_ns
]
records.sort()
assert len(records) == 2553, f"expected 2553 L1 pristine records, found {len(records)}"

codec = shutil.which("zstd")
assert codec, "zstd CLI is not available"
TEMP.mkdir()
per_lod = {}
try:
    for source in records:
        record = source.read_bytes()
        assert record[:8] == b"VFTILE16"
        raw_size, compressed_size = struct.unpack_from("<QQ", record, 124)
        raw = zlib.decompress(record[144:144 + compressed_size])
        assert len(raw) == raw_size
        lod = int(source.name.split("_", 1)[0][3:])
        entry = per_lod.setdefault(lod, {"raw": 0, "files": [], "encode_ns": 0, "decode_ns": 0, "zstd": 0})
        raw_file = TEMP / (source.stem + ".raw")
        raw_file.write_bytes(raw)
        entry["raw"] += len(raw)
        entry["files"].append(raw_file)

    for lod in sorted(per_lod):
        entry = per_lod[lod]
        start = time.perf_counter_ns()
        subprocess.run([codec, "-q", "-1", "-f", "--", *map(str, entry["files"])], check=True)
        entry["encode_ns"] = time.perf_counter_ns() - start
        compressed_files = [Path(str(path) + ".zst") for path in entry["files"]]
        entry["zstd"] = sum(path.stat().st_size for path in compressed_files)
        start = time.perf_counter_ns()
        subprocess.run([codec, "-q", "-d", "-f", "--", *map(str, compressed_files)], check=True)
        entry["decode_ns"] = time.perf_counter_ns() - start

    # Verify each decoded output byte-for-byte against the zlib payload again.
    # The filename is the tile stem, so a deterministic source lookup is sufficient.
    source_by_stem = {path.stem: path for path in records}
    for entry in per_lod.values():
        for raw_file in entry["files"]:
            source = source_by_stem[raw_file.stem]
            data = source.read_bytes()
            _, compressed_size = struct.unpack_from("<QQ", data, 124)
            assert raw_file.read_bytes() == zlib.decompress(data[144:144 + compressed_size])

    lines = ["lod tiles raw_total_B zstd1_total_B zstd1_avg_B encode_us_per_tile decode_us_per_tile status"]
    for lod in sorted(per_lod):
        entry = per_lod[lod]
        count = len(entry["files"])
        lines.append(
            f"{lod} {count} {entry['raw']} {entry['zstd']} {entry['zstd'] / count:.1f} "
            f"{entry['encode_ns'] / count / 1000:.2f} {entry['decode_ns'] / count / 1000:.2f} PASS"
        )
    raw_total = sum(entry["raw"] for entry in per_lod.values())
    zstd_total = sum(entry["zstd"] for entry in per_lod.values())
    lines.append(f"total_tiles={len(records)} raw_bytes={raw_total} zstd1_bytes={zstd_total} ratio={zstd_total / raw_total:.6f}")
    lines.append("decode=byte-exact; timing includes zstd CLI process startup once per LOD")
    RESULT.write_text("\n".join(lines) + "\n")
    print("\n".join(lines))
finally:
    resolved = TEMP.resolve()
    assert resolved.parent == CACHE.resolve()
    shutil.rmtree(TEMP)