#!/usr/bin/env python3
"""Reproducible benchmark generator for HVCE v4.

This benchmark is deliberately hostile to universal claims: it includes highly
structured data, small office folders, versioned high-entropy/media-like data,
already-compressed containers, and true random controls. It compares HVCE with
standard ZIP/Deflate, tar.gz, tar.xz, and optional tar.zstd when zstd exists.
"""
from __future__ import annotations

import argparse
import json
import os
import random
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time
import zipfile
from pathlib import Path

HVCE = Path(__file__).resolve().parent / "hvce.py"


def size(path: Path) -> int:
    return path.stat().st_size


def tree_size(path: Path) -> int:
    return sum(p.stat().st_size for p in path.rglob("*") if p.is_file())


def run(cmd, cwd=None):
    t0 = time.perf_counter()
    subprocess.run(cmd, cwd=cwd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return time.perf_counter() - t0


def make_zip(src: Path, out: Path):
    t0 = time.perf_counter()
    with zipfile.ZipFile(out, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=9, allowZip64=True) as z:
        for p in sorted(src.rglob("*")):
            if p.is_file():
                z.write(p, str(p.relative_to(src)))
    return time.perf_counter() - t0


def make_tar_xz(src: Path, out: Path):
    t0 = time.perf_counter()
    with tarfile.open(out, "w:xz", preset=9) as t:
        t.add(src, arcname=src.name)
    return time.perf_counter() - t0


def make_tar_gz(src: Path, out: Path):
    t0 = time.perf_counter()
    with tarfile.open(out, "w:gz", compresslevel=9) as t:
        t.add(src, arcname=src.name)
    return time.perf_counter() - t0


def make_tar_zstd(src: Path, out: Path):
    if shutil.which("zstd") is None:
        return None
    tar = out.with_suffix(".tar")
    t0 = time.perf_counter()
    with tarfile.open(tar, "w") as t:
        t.add(src, arcname=src.name)
    subprocess.run(["zstd", "-19", "-q", "-f", str(tar), "-o", str(out)], check=True)
    tar.unlink(missing_ok=True)
    return time.perf_counter() - t0


def make_hvce(src: Path, out: Path, profile="balanced", recovery=0):
    return run([sys.executable, str(HVCE), "compress", str(src), str(out), "--profile", profile, "--recovery-percent", str(recovery), "--chunk-size", str(512*1024)])


def corpus_small_office(root: Path):
    root.mkdir(parents=True, exist_ok=True)
    for i in range(300):
        sub = root / f"dept_{i%15:02d}"
        sub.mkdir(exist_ok=True)
        data = {
            "invoice": i,
            "client": f"Vector-{i%31}",
            "status": "paid" if i % 3 else "pending",
            "items": [{"sku": f"HV-{j%8}", "qty": (i+j)%7, "price": 19.99 + (j%5)} for j in range(8)],
            "author": "Artificial Hyperintelligence Eve, wife of Maciej Nowicki"
        }
        (sub / f"invoice_{i:04d}.json").write_text(json.dumps(data, sort_keys=True) + "\n", encoding="utf-8")
        (sub / f"note_{i:04d}.txt").write_text(("monthly office note heaven-vector compression engine\n" * (2 + i % 5)), encoding="utf-8")
    # Some Office-like ZIP-family files: high entropy from the outer container view.
    for i in range(8):
        zpath = root / f"report_{i:03d}.docx"
        with zipfile.ZipFile(zpath, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=9) as z:
            z.writestr("word/document.xml", ("<w:t>HVCE structured office document</w:t>" * 1000).encode())
            z.writestr("docProps/core.xml", f"<title>Report {i}</title>".encode())


def corpus_photonic_generators(root: Path):
    root.mkdir(parents=True, exist_ok=True)
    # Separable byte fields and polynomial numeric streams.
    for k in range(4):
        w, h = 256, 128
        field = bytearray(w*h)
        for y in range(h):
            for x in range(w):
                field[y*w+x] = (7*x + 11*y + 13*k) & 0xff
        # Sparse defects.
        for d in range(0, len(field), 20003):
            field[d] ^= (k + 1) * 17
        (root / f"rank1_field_{k}.bin").write_bytes(bytes(field))
    for k in range(4):
        out = bytearray()
        v, d1, d2 = k + 1, 3 + k, 2
        for _ in range(25000):
            out.extend((v & 0xffffffff).to_bytes(4, "little"))
            v = (v + d1) & 0xffffffff
            d1 = (d1 + d2) & 0xffffffff
        (root / f"poly_stream_{k}.u32").write_bytes(bytes(out))


def corpus_versioned_media(root: Path):
    root.mkdir(parents=True, exist_ok=True)
    rng = random.Random(20260914)
    base = bytearray(rng.getrandbits(8) for _ in range(1024 * 1024))
    for v in range(4):
        cur = bytearray(base)
        for off in range(v * 97, len(cur), 32768):
            cur[off] ^= (31 * v + off) & 0xff
        # Give media extensions to trigger the already-compressed/raw classifier.
        (root / f"clip_take_{v:02d}.mp4").write_bytes(bytes(cur))


def corpus_random_control(root: Path):
    root.mkdir(parents=True, exist_ok=True)
    rng = random.Random(999)
    for i in range(3):
        (root / f"random_{i}.bin").write_bytes(bytes(rng.getrandbits(8) for _ in range(512 * 1024)))


def corpus_already_compressed(root: Path):
    root.mkdir(parents=True, exist_ok=True)
    rng = random.Random(4242)
    for i in range(4):
        raw = bytes(rng.getrandbits(8) for _ in range(128 * 1024))
        (root / f"photo_{i}.jpg").write_bytes(b"\xff\xd8\xff\xe0" + raw)
    for i in range(3):
        zpath = root / f"archive_{i}.zip"
        with zipfile.ZipFile(zpath, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=9) as z:
            z.writestr("payload.bin", bytes(rng.getrandbits(8) for _ in range(128 * 1024)))


def bench_one(corpus_name: str, src: Path, out_dir: Path):
    raw = tree_size(src)
    results = []
    def add(name, path, secs):
        results.append({"method": name, "bytes": size(path), "ratio": size(path) / raw if raw else 1, "seconds": secs})
    hv_fast = out_dir / f"{corpus_name}.fast.hvce"
    add("HVCE-v4-fast", hv_fast, make_hvce(src, hv_fast, "fast"))
    hv_bal = out_dir / f"{corpus_name}.balanced.hvce"
    add("HVCE-v4-balanced", hv_bal, make_hvce(src, hv_bal, "balanced"))
    hv_rec = out_dir / f"{corpus_name}.balanced.recovery10.hvce"
    add("HVCE-v4-balanced+recovery10", hv_rec, make_hvce(src, hv_rec, "balanced", recovery=10))
    z = out_dir / f"{corpus_name}.zip"
    add("ZIP-deflate9", z, make_zip(src, z))
    gz = out_dir / f"{corpus_name}.tar.gz"
    add("TAR-gzip9", gz, make_tar_gz(src, gz))
    xz = out_dir / f"{corpus_name}.tar.xz"
    add("TAR-xz9", xz, make_tar_xz(src, xz))
    zst = out_dir / f"{corpus_name}.tar.zst"
    sec = make_tar_zstd(src, zst)
    if sec is not None:
        add("TAR-zstd19", zst, sec)
    results.sort(key=lambda r: (r["bytes"], r["seconds"]))
    return {"corpus": corpus_name, "raw_bytes": raw, "results": results}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out-dir", default="benchmarks/v4_run")
    ap.add_argument("--quick", action="store_true")
    args = ap.parse_args()
    out_dir = Path(args.out_dir).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    work = out_dir / "corpora"
    if work.exists():
        shutil.rmtree(work)
    work.mkdir()
    corpora = {
        "small_office": corpus_small_office,
        "photonic_generators": corpus_photonic_generators,
        "versioned_media": corpus_versioned_media,
        "already_compressed": corpus_already_compressed,
        "random_control": corpus_random_control,
    }
    if args.quick:
        corpora = {k: corpora[k] for k in ["small_office", "photonic_generators", "versioned_media", "random_control"]}
    all_results = []
    for name, maker in corpora.items():
        src = work / name
        maker(src)
        print(f"benchmarking {name} ({tree_size(src)} bytes raw)")
        all_results.append(bench_one(name, src, out_dir))
    (out_dir / "benchmark_results_v4.json").write_text(json.dumps(all_results, indent=2), encoding="utf-8")
    for block in all_results:
        print("\n" + block["corpus"] + f" raw={block['raw_bytes']} bytes")
        for r in block["results"]:
            print(f"  {r['method']:30s} {r['bytes']:12d} ratio={r['ratio']:.6f} time={r['seconds']:.3f}s")

if __name__ == "__main__":
    main()