#!/usr/bin/env python3
"""Compare token cost of equivalent YAML and JSON documents.

Run:
    cd scripts/benchmarks/yaml-vs-json-tokens
    python3 -m venv venv
    ./venv/bin/pip install tiktoken
    ./venv/bin/python measure.py
"""
from pathlib import Path
import json
import tiktoken

HERE = Path(__file__).parent
ENCODINGS = ["cl100k_base", "o200k_base"]


def main() -> None:
    encoders = {name: tiktoken.get_encoding(name) for name in ENCODINGS}
    files = {
        "JSON (pretty)": HERE / "sample.json",
        "YAML": HERE / "sample.yaml",
    }
    rows = {}
    for label, path in files.items():
        text = path.read_text(encoding="utf-8")
        rows[label] = {
            "chars": len(text),
            "bytes": len(text.encode("utf-8")),
            "cl100k": len(encoders["cl100k_base"].encode(text)),
            "o200k": len(encoders["o200k_base"].encode(text)),
        }

    # Also measure compact (single-line) JSON for fairness, since some
    # callers strip whitespace before sending to the model.
    pretty_json = (HERE / "sample.json").read_text(encoding="utf-8")
    compact_json = json.dumps(json.loads(pretty_json), separators=(",", ":"))
    rows["JSON (compact)"] = {
        "chars": len(compact_json),
        "bytes": len(compact_json.encode("utf-8")),
        "cl100k": len(encoders["cl100k_base"].encode(compact_json)),
        "o200k": len(encoders["o200k_base"].encode(compact_json)),
    }

    print(f"tiktoken {tiktoken.__version__}")
    header = f"{'format':<18}{'chars':>8}{'bytes':>8}{'cl100k':>10}{'o200k':>10}"
    print(header)
    print("-" * len(header))
    for label, m in rows.items():
        print(f"{label:<18}{m['chars']:>8}{m['bytes']:>8}{m['cl100k']:>10}{m['o200k']:>10}")

    # Reduction summary
    print()
    json_pretty = rows["JSON (pretty)"]
    yaml_ = rows["YAML"]
    json_compact = rows["JSON (compact)"]
    for ref_label, ref in [("JSON (pretty)", json_pretty), ("JSON (compact)", json_compact)]:
        diff_chars = (yaml_["chars"] - ref["chars"]) / ref["chars"] * 100
        diff_cl100k = (yaml_["cl100k"] - ref["cl100k"]) / ref["cl100k"] * 100
        diff_o200k = (yaml_["o200k"] - ref["o200k"]) / ref["o200k"] * 100
        print(
            f"YAML vs {ref_label}: chars {diff_chars:+.1f}%, "
            f"cl100k {diff_cl100k:+.1f}%, o200k {diff_o200k:+.1f}%"
        )

    import sys
    out = {
        "topic": "YAML vs JSON size and token cost (same 32-line web/K8s-style sample)",
        "environment": {
            "machine": "Apple M5 Pro (macOS, Darwin 25.5.0)",
            "python": sys.version.split()[0],
            "tiktoken": tiktoken.__version__,
        },
        "rows": rows,
    }
    (HERE / "results.json").write_text(json.dumps(out, indent=2) + "\n", encoding="utf-8")
    print("results.json written")


if __name__ == "__main__":
    main()