# /// script
# requires-python = ">=3.12,<3.13"
# dependencies = ["tiktoken==0.14.0", "toon-format==1.1.0"]
# ///
"""Recount checked-in synthetic payloads; never contact a model or a live account.

Run: uv run scripts/measure-mcp-context.py [--write]
JSON baseline for encoding-only comparisons is minified, UTF-8, same data as TOON.
Sequence counts tokenize serialized request text, not provider chat framing or bills.
"""
import argparse
import json
from pathlib import Path

import tiktoken
from toon_format import dumps, loads

ROOT = Path(__file__).resolve().parents[1]
FIXTURES = ROOT / "public/blog/evidence/mcp-context-fixtures-2026-10-06.json"
RESULTS = ROOT / "src/data/mcp-context-measurements.json"
PUBLIC_RESULTS = ROOT / "public/blog/evidence/mcp-context-measurements-2026-10-06.json"


def compact(value):
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def reduction(before, after):
    return round(100 * (1 - after / before), 2)


def json_responses(raw):
    request = json.loads(raw)
    for message in request["messages"]:
        if message.get("role") == "tool" and message.get("tool_call_id") in ("search", "definition"):
            message["content"] = compact(loads(message["content"]))
    return compact(request)


def measure(fixtures_path=FIXTURES):
    fixtures = json.loads(fixtures_path.read_text())
    report = {key: fixtures[key] for key in ("date", "scope", "source_revisions", "runtime", "capture", "packages", "sample")}
    report["encoding_baseline"] = "Minified JSON; exactly the decoded current TOON value, same fields, constraints and order."
    report["sequence_baseline"] = "Previous verbose discovery JSON (original default json.dumps spacing) and helper contract; current captured service response and compact helpers. Repeated history and an exact-schema lookup, when required, are included."
    report["limitations"] = "Deterministic synthetic sample, one selected action, fixed user prompt and fixed successful result. No additional system prompt or pre-existing conversation. No model choices, retries, quality evaluation, provider message framing, output generation, latency, cache treatment or billed cost. Percentages have different denominators and must not be added."
    report["tokenizers"] = {}
    for tokenizer in ("o200k_base", "cl100k_base"):
        encoding = tiktoken.get_encoding(tokenizer)
        count = lambda text: len(encoding.encode(text))
        cases = []
        for sample in fixtures["samples"]:
            text = sample["current_response"]
            value = loads(text)
            assert loads(dumps(value)) == value, sample["id"]
            # Schemas retained in discovery must be byte-for-value identical to native input.
            legacy = json.loads(sample["legacy_response"])
            schemas = {tool["display_name"]: tool["inputSchema"] for tool in legacy["tools"]}
            for tool in value.get("tools", [value]):
                if "inputSchema" in tool:
                    assert tool["inputSchema"] == schemas[tool["name"].split(".", 1)[1]]
            before = count(sample["legacy_response"])
            projected = count(compact(value))
            after = count(text)
            row = {key: sample[key] for key in ("id", "shape", "candidates", "mode")}
            row["schema_in_search"] = "inputSchema" in text
            row["response"] = {"legacy_json": before, "projected_json": projected, "toon": after,
                "combined_reduction_percent": reduction(before, after),
                "encoding_reduction_percent": reduction(projected, after)}
            if "legacy_requests" in sample:
                old = sum(map(count, sample["legacy_requests"]))
                new = sum(map(count, sample["current_requests"]))
                as_json = sum(count(json_responses(raw)) for raw in sample["current_requests"])
                row["sequence"] = {"legacy": old, "current_json": as_json, "current_toon": new,
                    "legacy_rounds": len(sample["legacy_requests"]), "current_rounds": len(sample["current_requests"]),
                    "combined_reduction_percent": reduction(old, new), "encoding_reduction_percent": reduction(as_json, new)}
            cases.append(row)
        initial = []
        for sample in fixtures["initial_context"]:
            eager, discovery = count(sample["eager_request"]), count(sample["discovery_request"])
            initial.append({"shape": sample["shape"], "catalog_tools": sample["catalog_tools"],
                "eager": eager, "discovery": discovery, "reduction_percent": reduction(eager, discovery)})
        report["tokenizers"][tokenizer] = {"cases": cases, "initial_context": initial}
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--write", action="store_true", help="Regenerate the results after reviewing fixture changes.")
    parser.add_argument("--fixtures", type=Path, default=FIXTURES, help="Path to the captured fixture JSON.")
    parser.add_argument("--results", type=Path, default=RESULTS, help="Path to the measurement JSON to verify or regenerate.")
    args = parser.parse_args()
    report = measure(args.fixtures)
    if args.write:
        content = json.dumps(report, ensure_ascii=False, indent=2) + "\n"
        args.results.write_text(content)
        if args.results == RESULTS:
            PUBLIC_RESULTS.write_text(content)
    else:
        assert json.loads(args.results.read_text()) == report, "Measurement artifact differs; inspect fixtures and recount."
        if args.results == RESULTS:
            assert PUBLIC_RESULTS.read_text() == RESULTS.read_text(), "Public results differ from the displayed dataset."
    for tokenizer, values in report["tokenizers"].items():
        print(tokenizer)
        for row in values["cases"]:
            response = row["response"]
            sequence = row.get("sequence")
            print(row["id"], f'encoding: {response["projected_json"]} -> {response["toon"]} ({response["encoding_reduction_percent"]}%)',
                f'sequence: {sequence["legacy"]} -> {sequence["current_toon"]} ({sequence["combined_reduction_percent"]}%)' if sequence else '')
        for row in values["initial_context"]:
            print('initial', row)
    print("Verified 11 same-data response pairs and 10 deterministic sequences across two tokenizers.")


if __name__ == "__main__":
    main()