#!/usr/bin/env python3
"""Measure actual SeaweedFS S3 objects; keep training and evaluation disjoint."""
import argparse
import collections
import csv
import datetime
import gzip
import hashlib
import html
import json
import math
import pathlib
import platform
import statistics
import subprocess
import time
import urllib.request

import boto3
from botocore.config import Config

ARCHIVE = "2025-01-01-0.json.gz"
ARCHIVE_SHA256 = "477e6ab5ab27e7217a4807f9836ef8e947983f9637b26311e2c34432dd3949f8"
S3 = "http://127.0.0.1:8533"
FILER = "http://127.0.0.1:9088"
MASTER = "127.0.0.1:9533"
BUCKETS = {"plain": "gh-events-plain", "dict": "gh-events-dict"}


def digest(data):
    return hashlib.sha256(data).hexdigest()


def client():
    return boto3.client("s3", endpoint_url=S3, region_name="us-east-1",
                        aws_access_key_id="blog", aws_secret_access_key="blog-secret",
                        config=Config(s3={"addressing_style": "path"}, retries={"max_attempts": 0}))


def get(url, headers=None):
    with urllib.request.urlopen(urllib.request.Request(url, headers=headers or {}), timeout=30) as r:
        return r.read(), dict(r.headers)


def metadata(path):
    return json.loads(get(FILER + path + "?metadata=true")[0])


def raw_object(bucket, key):
    meta = metadata(f"/buckets/{bucket}/{key}")
    chunks = meta.get("chunks", [])
    assert len(chunks) == 1, (key, meta)
    fid = chunks[0]["file_id"]
    lookup = json.loads(get(f"http://{MASTER}/dir/lookup?volumeId={fid.split(',')[0]}")[0])
    body, headers = get("http://" + lookup["locations"][0]["url"] + "/" + fid,
                        {"Accept-Encoding": "zstd", "x-seaweed-zstd-dict-accept": "1"})
    return body, {k.lower(): v for k, v in headers.items()}


def records(root, split):
    return [p.read_bytes() for p in sorted((root / "data" / split).glob("*.json"))]


def shell(root, command, filename):
    proc = subprocess.run([str(root / "weed"), "shell", "-master=" + MASTER,
                           "-filer=127.0.0.1:9088"], input=command + "\n",
                          text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=True)
    text = "$ " + command + "\n" + proc.stdout
    (root / "evidence" / filename).write_text(text)
    print(text, end="", flush=True)
    return text


def prepare(root):
    archive = root / "data" / ARCHIVE
    assert digest(archive.read_bytes()) == ARCHIVE_SHA256, "Unexpected archive bytes"
    selected = []
    with gzip.open(archive, "rb") as source:
        for i, line in enumerate(source):
            if i == 2000:
                break
            json.loads(line)  # Validate, preserving the original bytes and newline.
            split = "train" if i < 1000 else "eval"
            directory = root / "data" / split
            directory.mkdir(parents=True, exist_ok=True)
            (directory / f"{i:04d}.json").write_bytes(line)
            selected.append(line)
    assert len(selected) == 2000
    ids = [json.loads(row)["id"] for row in selected]
    assert len(set(ids)) == 2000, "Duplicate event IDs across the split"
    info = {"source": "https://data.gharchive.org/" + ARCHIVE,
            "archive_sha256": ARCHIVE_SHA256, "train_records": 1000, "eval_records": 1000,
            "train_sha256": digest(b"".join(selected[:1000])),
            "eval_sha256": digest(b"".join(selected[1000:])),
            "eval_logical_bytes": sum(map(len, selected[1000:])),
            "eval_min_bytes": min(map(len, selected[1000:])),
            "eval_max_bytes": max(map(len, selected[1000:])),
            "eval_median_bytes": statistics.median(map(len, selected[1000:]))}
    (root / "evidence" / "dataset.json").write_text(json.dumps(info, indent=2) + "\n")
    print(json.dumps(info, indent=2), flush=True)


def seed(root):
    s3 = client()
    for bucket in BUCKETS.values():
        s3.create_bucket(Bucket=bucket)
    for i, row in enumerate(records(root, "train")):
        s3.put_object(Bucket=BUCKETS["dict"], Key=f"train/{i:04d}.json", Body=row,
                      ContentType="application/json")
    print("Uploaded 1000 original training records; evaluation records have not been uploaded.", flush=True)
    shell(root, "fs.zstd.dict -bucket gh-events-dict -samples 512", "dry-run.txt")
    capture = shell(root, "fs.zstd.dict -bucket gh-events-dict -samples 512 -apply", "training.txt")
    assert "bound to /etc/zstd/dicts/" in capture
    shell(root, "fs.zstd.dict -bucket gh-events-dict -show", "binding.txt")
    print("Waiting 31 seconds for the writer's 30-second binding cache.", flush=True)
    time.sleep(31)


def timings(s3, rows, operation):
    values = {mode: [] for mode in BUCKETS}
    # Warm both paths; alternate which path runs first to reduce ordering bias.
    for i, row in enumerate(rows[:350]):
        for mode in (list(BUCKETS) if i % 2 == 0 else list(reversed(BUCKETS))):
            args = dict(Bucket=BUCKETS[mode], Key=f"latency/{i:04d}.json")
            start = time.perf_counter_ns()
            if operation == "put":
                s3.put_object(**args, Body=row, ContentType="application/json")
            else:
                assert s3.get_object(**args)["Body"].read() == row
            elapsed = (time.perf_counter_ns() - start) / 1000
            if i >= 50:
                values[mode].append(elapsed)
    return {mode: {"requests": len(v), "median_us": round(statistics.median(v), 1),
                   "p95_us": round(sorted(v)[math.ceil(len(v) * .95) - 1], 1)}
            for mode, v in values.items()}


def measure(root):
    s3 = client()
    rows = records(root, "eval")
    measured, seen = [], {mode: collections.Counter() for mode in BUCKETS}
    for i, row in enumerate(rows):
        key = f"eval/{i:04d}.json"
        record = json.loads(row)
        result = {"event_id": record["id"], "type": record["type"], "logical_bytes": len(row)}
        for mode, bucket in BUCKETS.items():
            s3.put_object(Bucket=bucket, Key=key, Body=row, ContentType="application/json")
            assert s3.get_object(Bucket=bucket, Key=key)["Body"].read() == row
            raw, headers = raw_object(bucket, key)
            assert headers.get("content-encoding") == "zstd", (mode, key, headers)
            assert raw.startswith(bytes.fromhex("28b52ffd")), (mode, key)
            dict_id = int(headers.get("x-seaweed-zstd-dict-id", "0"))
            assert (dict_id != 0) == (mode == "dict"), (mode, key, dict_id)
            seen[mode][dict_id] += 1
            result[mode + "_bytes"] = len(raw)
            result[mode + "_dict_id"] = dict_id
            if i == 0:
                (root / "evidence" / (mode + "-sample.zst")).write_bytes(raw)
        measured.append(result)
        if (i + 1) % 100 == 0:
            print(f"Measured and verified {i + 1}/1000 evaluation pairs.", flush=True)
    with (root / "evidence" / "objects.csv").open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(measured[0]), lineterminator="\n")
        writer.writeheader()
        writer.writerows(measured)

    binding = metadata("/buckets/" + BUCKETS["dict"])
    # Registry is isolated to this experiment, so count every registered dictionary.
    listing = json.loads(get(FILER + "/etc/zstd/dicts/?limit=100", {"Accept": "application/json"})[0])
    dictionaries = []
    for entry in listing["Entries"]:
        name = entry["FullPath"].rsplit("/", 1)[-1]
        body = get(FILER + "/etc/zstd/dicts/" + name)[0]
        (root / "evidence" / name).write_bytes(body)
        dictionaries.append({"name": name, "bytes": len(body), "sha256": digest(body)})
    logical = sum(x["logical_bytes"] for x in measured)
    plain = sum(x["plain_bytes"] for x in measured)
    compressed = sum(x["dict_bytes"] for x in measured)
    overhead = sum(x["bytes"] for x in dictionaries)
    per_type = []
    for kind in sorted({x["type"] for x in measured}):
        group = [x for x in measured if x["type"] == kind]
        per_type.append({"type": kind, "count": len(group),
                         "logical_bytes": sum(x["logical_bytes"] for x in group),
                         "plain_bytes": sum(x["plain_bytes"] for x in group),
                         "dict_bytes": sum(x["dict_bytes"] for x in group)})
    # Extension control: JSON MIME compression and extension/default dictionary selection.
    s3.put_object(Bucket=BUCKETS["dict"], Key="control/application-json.json",
                  Body=rows[0], ContentType="application/json")
    control, control_headers = raw_object(BUCKETS["dict"], "control/application-json.json")
    assert control_headers.get("content-encoding") == "zstd" and len(control) < len(rows[0])
    s3.put_object(Bucket=BUCKETS["dict"], Key="control/application-json.bin",
                  Body=rows[0], ContentType="application/json")
    opaque, opaque_headers = raw_object(BUCKETS["dict"], "control/application-json.bin")
    assert opaque_headers.get("content-encoding") == "zstd" and len(opaque) < len(rows[0])
    assert control_headers["x-seaweed-zstd-dict-id"] != opaque_headers["x-seaweed-zstd-dict-id"]
    result = {"measured_at_utc": datetime.datetime.now(datetime.timezone.utc).isoformat(),
              "dataset": json.loads((root / "evidence" / "dataset.json").read_text()),
              "environment": {"host": "swdev", "kernel": platform.release(),
                              "python": platform.python_version(), "boto3": boto3.__version__,
                              "compression_level": 1, "client_concurrency": 1,
                              "source_commit": (root / "source-commit.txt").read_text().strip()},
              "logical_bytes": logical, "plain_bytes": plain, "dict_bytes": compressed,
              "dictionary_registry_bytes": overhead, "dictionaries": dictionaries,
              "payload_saving_percent": round(100 * (plain - compressed) / plain, 2),
              "saving_including_registry_percent": round(100 * (plain - compressed - overhead) / plain, 2),
              "break_even_objects": math.ceil(overhead / ((plain - compressed) / len(rows))),
              "frame_dictionary_ids": {m: dict(v) for m, v in seen.items()},
              "per_type": per_type, "sample": measured[0],
              "content_type_control": {"type": "application/json", "logical_bytes": len(rows[0]),
                                       "json_key_stored_bytes": len(control), "json_key_compressed": True,
                                       "bin_key_stored_bytes": len(opaque), "bin_key_compressed": True,
                                       "json_dict_id": int(control_headers["x-seaweed-zstd-dict-id"]),
                                       "bin_dict_id": int(opaque_headers["x-seaweed-zstd-dict-id"])},
              "latency": {op: timings(s3, rows, op) for op in ("put", "get")}}
    (root / "evidence" / "binding-metadata.json").write_text(json.dumps(binding, indent=2) + "\n")
    shell(root, "fs.zstd.dict -bucket gh-events-dict -unbind -apply", "unbind.txt")
    for i, row in enumerate(rows):
        assert s3.get_object(Bucket=BUCKETS["dict"], Key=f"eval/{i:04d}.json")["Body"].read() == row
    verification = ("1000/1000 plain-zstd S3 GETs match the original archive bytes.\n"
                    "1000/1000 dictionary-zstd S3 GETs match the original archive bytes.\n"
                    "1000/1000 dictionary objects still match after unbinding.\n"
                    "1000/1000 dictionary frames name a nonzero dictionary ID.\n"
                    "1000/1000 plain frames use no dictionary.\n"
                    "application/json with a .json key is compressed with a dictionary.\n"
                    "application/json with a .bin key uses the bucket default dictionary.\n")
    (root / "evidence" / "verification.txt").write_text(verification)
    (root / "evidence" / "results.json").write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2), flush=True)
    print(verification, end="", flush=True)


def report(root):
    evidence = root / "evidence"
    r = json.loads((evidence / "results.json").read_text())
    def number(n):
        return f"{n:,}"
    def capture(name):
        return html.escape((evidence / name).read_text())
    types = "".join(f"<tr><td>{x['type']}</td><td>{x['count']}</td><td>{number(x['plain_bytes'])}</td>"
                    f"<td>{number(x['dict_bytes'])}</td><td>{100*(1-x['dict_bytes']/x['plain_bytes']):.1f}%</td></tr>"
                    for x in r["per_type"])
    page = f"""<!doctype html><html lang="en"><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>Zstd dictionary experiment · swdev</title>
<style>body{{margin:0;background:#f1f5f9;color:#0f172a;font:17px/1.55 system-ui,sans-serif}}
main{{max-width:1140px;margin:32px auto;padding:32px;background:white;border:1px solid #e2e8f0;border-radius:12px}}
h1{{font-size:32px;margin:4px 0 10px}}h2{{font-size:23px;margin:30px 0 12px}}p{{margin:8px 0 16px}}
.eyebrow{{color:#2563eb;font-weight:650}}.cards{{display:flex;gap:18px;margin:24px 0}}
.card{{flex:1;background:#eff6ff;border:1px solid #bfdbfe;border-radius:8px;padding:18px}}
.card b{{display:block;font-size:30px;color:#1d4ed8}}table{{width:100%;border-collapse:collapse}}
th,td{{padding:10px 13px;text-align:right;border-bottom:1px solid #e2e8f0}}th:first-child,td:first-child{{text-align:left}}
th{{background:#f8fafc}}pre{{background:#0f172a;color:#e2e8f0;border-radius:8px;padding:20px;overflow:auto;font:14px/1.65 ui-monospace,monospace}}
.note{{color:#475569;font-size:15px}}a{{color:#2563eb}}code{{font-family:ui-monospace,monospace}}</style>
<main><div class="eyebrow">EXPERIMENT REPORT · swdev · {html.escape(r['measured_at_utc'][:10])}</div>
<h1>Custom zstd dictionaries on real GitHub events</h1>
<p>1,000 untouched evaluation records, one JSON event per S3 object. Separate first 1,000 records for training.<br>
GH Archive: 2025-01-01 00:00 UTC · SeaweedFS level 1 · identical input bytes on both paths.</p>
<div class="cards"><div class="card"><b>{r['payload_saving_percent']:.1f}%</b>smaller than plain zstd</div>
<div class="card"><b>{r['saving_including_registry_percent']:.1f}%</b>smaller including both dictionaries</div>
<div class="card"><b>3,000 / 3,000</b>byte-for-byte read checks passed</div></div>
<table><thead><tr><th>Storage representation</th><th>Bytes</th><th>Relative to original</th></tr></thead><tbody>
<tr><td>Original JSON</td><td>{number(r['logical_bytes'])}</td><td>100.0%</td></tr>
<tr><td>Plain zstd</td><td>{number(r['plain_bytes'])}</td><td>{100*r['plain_bytes']/r['logical_bytes']:.1f}%</td></tr>
<tr><td>Zstd + dictionary (object payloads)</td><td>{number(r['dict_bytes'])}</td><td>{100*r['dict_bytes']/r['logical_bytes']:.1f}%</td></tr>
<tr><td>Dictionary payloads + entire registry</td><td>{number(r['dict_bytes']+r['dictionary_registry_bytes'])}</td><td>{100*(r['dict_bytes']+r['dictionary_registry_bytes'])/r['logical_bytes']:.1f}%</td></tr>
</tbody></table><p class="note">Measured raw volume payload bytes. Excludes needle/index overhead, filer metadata, replication and filesystem allocation.
Both registered dictionaries total {number(r['dictionary_registry_bytes'])} bytes; payload break-even is {r['break_even_objects']} similar objects.</p>
<h2>Different event types, different savings</h2><table><thead><tr><th>Event type</th><th>Objects</th><th>Plain zstd bytes</th><th>Dictionary bytes</th><th>Saving</th></tr></thead><tbody>{types}</tbody></table>
<p class="note">This is a report generated by the experiment script, separate from the SeaweedFS Admin UI.
<a href="results.json">Raw results</a> · <a href="objects.csv">Per-object CSV</a></p></main>
<main id="training"><div class="eyebrow">ACTUAL WEED SHELL CAPTURE</div><h1>Train, compare held-out samples, bind</h1>
<p>Only the training records existed in this bucket when the command ran. The final 1,000 evaluation records were still held back.</p>
<pre>{capture('training.txt')}</pre><h2>Published binding</h2><pre>{capture('binding.txt')}</pre></main>
<main id="verification"><div class="eyebrow">READ AND CONTENT-TYPE CONTROLS</div><h1>Same content through S3, including after unbind</h1>
<pre>{capture('verification.txt')}</pre><h2>Stop using the dictionary for new writes</h2><pre>{capture('unbind.txt')}</pre>
<p>Uploads use <code>ContentType='application/json'</code> and <code>.json</code> object keys.
The .json control stores {r['content_type_control']['json_key_stored_bytes']:,} compressed bytes;
the same content under a .bin key uses the default dictionary and stores {r['content_type_control']['bin_key_stored_bytes']:,} compressed bytes.</p></main></html>"""
    (evidence / "report.html").write_text(page)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("phase", choices=["prepare", "seed", "measure", "report"])
    parser.add_argument("--root", type=pathlib.Path, required=True)
    args = parser.parse_args()
    globals()[args.phase](args.root.resolve())