#!/usr/bin/env python3

"""Calibrate char/4 token estimates against real tokenizer behavior.

The benchmark harnesses estimate tokens as characters divided by four. That is

biased for VCS output: hex hashes tokenize at roughly 2-3 chars/token and

diff/JSON syntax tokenizes worse than prose, so char/4 UNDERESTIMATES real

tokens — and by a different factor per subject depending on output style.

Comparing Oak vs Git token deltas without knowing each side's bias risks

publishing the wrong number.

This tool reports, per input sample:

- chars and the naive char/4 estimate

- a composition-calibrated estimate from documented per-class rates

- the exact token count when a real tokenizer is importable

  (`tiktoken`, optional; never required)

- the calibration factor relative to char/4

Feed it captured command output (for example agent-workflow transcript

artifacts under results/agent-workflow/artifacts, or any saved stdout).

Group files with --label to get per-subject or per-operation factors.

Usage:

    python3 scripts/token_calibration.py --self-test

    python3 scripts/token_calibration.py path/to/outputs/*.log --label oak_diff

    python3 scripts/token_calibration.py results/agent-workflow/artifacts --label claude_runs

"""

from __future__ import annotations

import argparse

import json

import re

import sys

from pathlib import Path

from typing import Any

# Approximate chars-per-token rates by character class, derived from how

# BPE-family tokenizers (cl100k/o200k and Claude's tokenizer behave similarly

# in aggregate) split common VCS output. These are documented approximations

# for bias *direction and rough magnitude*; use a real tokenizer for exact

# counts. Sources: hex strings split into 2-3 char fragments; punctuation

# rarely merges across symbols; prose with spaces averages ~4.4 chars/token.

CLASS_RATES = {

    "hex_run": 2.4,        # commit hashes, object ids, checksums

    "digits": 2.7,         # line numbers, counts, sizes

    "punctuation": 2.0,    # diff markers, JSON syntax, path separators

    "prose": 4.4,          # words plus the spaces that merge into them

}

HEX_RUN = re.compile(r"[0-9a-fA-F]{8,}")

DIGITS = re.compile(r"[0-9]")

PUNCT = re.compile(r"[^\w\s]")

WORD_OR_SPACE = re.compile(r"[\w\s]")

DEFAULT_MAX_BYTES = 5 * 1024 * 1024

SELF_TEST_SAMPLES = {

    "git_diff_like": (

        "diff --git a/app/pricing.py b/app/pricing.py\n"

        "index 3f9c2ab8e1d44c7a..b82f1c9d0e3a55f1 100644\n"

        "--- a/app/pricing.py\n"

        "+++ b/app/pricing.py\n"

        "@@ -1,4 +1,4 @@\n"

        " def apply_discount(cents: int, percent: int) -> int:\n"

        "-    return cents - percent\n"

        "+    return cents - ((cents * percent) // 100)\n"

    ) * 8,

    "hash_heavy_log": (

        "commit 4f2a9c8be1d34f7a9b82f1c9d0e3a55f19c8be1d\n"

        "commit a9b82f1c9d0e3a55f19c8be1d34f7a9b82f1c9d0\n"

        "commit 0e3a55f19c8be1d34f7a9b82f1c9d04f2a9c8be1\n"

    ) * 12,

    "prose_readme": (

        "The benchmark suite measures how much the version control model helps "

        "or hurts real coding agent workflows without baking in assumptions "

        "about a specific model provider or interface. "

    ) * 10,

}

def classify_chars(text: str) -> dict[str, int]:

    counts = {"hex_run": 0, "digits": 0, "punctuation": 0, "prose": 0}

    remainder_parts: list[str] = []

    last_end = 0

    for match in HEX_RUN.finditer(text):

        counts["hex_run"] += match.end() - match.start()

        remainder_parts.append(text[last_end : match.start()])

        last_end = match.end()

    remainder_parts.append(text[last_end:])

    remainder = "".join(remainder_parts)

    counts["digits"] = len(DIGITS.findall(remainder))

    counts["punctuation"] = len(PUNCT.findall(remainder))

    counts["prose"] = len(remainder) - counts["digits"] - counts["punctuation"]

    return counts

def calibrated_estimate(text: str) -> tuple[float, dict[str, int]]:

    counts = classify_chars(text)

    estimate = sum(counts[name] / CLASS_RATES[name] for name in counts)

    return estimate, counts

def exact_tokens(text: str) -> tuple[int | None, str | None]:

    try:

        import tiktoken  # type: ignore[import-not-found]

    except ImportError:

        return None, None

    for encoding_name in ("o200k_base", "cl100k_base"):

        try:

            encoding = tiktoken.get_encoding(encoding_name)

        except Exception:

            continue

        return len(encoding.encode(text, disallowed_special=())), encoding_name

    return None, None

def analyze_text(name: str, text: str) -> dict[str, Any]:

    chars = len(text)

    char_div_4 = max(1, (chars + 3) // 4) if chars else 0

    calibrated, composition = calibrated_estimate(text)

    exact, tokenizer = exact_tokens(text)

    best = exact if exact is not None else calibrated

    return {

        "sample": name,

        "chars": chars,

        "char_div_4_estimate": char_div_4,

        "composition": composition,

        "calibrated_estimate": round(calibrated, 1),

        "exact_tokens": exact,

        "tokenizer": tokenizer,

        "best_estimate": round(float(best), 1),

        "calibration_factor_vs_char_div_4": round(float(best) / char_div_4, 3) if char_div_4 else None,

    }

def iter_input_files(paths: list[Path]) -> list[Path]:

    files: list[Path] = []

    for path in paths:

        if path.is_dir():

            files.extend(sorted(p for p in path.rglob("*") if p.is_file()))

        elif path.is_file():

            files.append(path)

        else:

            raise SystemExit(f"Input not found: {path}")

    return files

def aggregate(results: list[dict[str, Any]], label: str | None) -> dict[str, Any]:

    total_chars = sum(int(item["chars"]) for item in results)

    total_naive = sum(int(item["char_div_4_estimate"]) for item in results)

    total_best = sum(float(item["best_estimate"]) for item in results)

    exact_count = sum(1 for item in results if item["exact_tokens"] is not None)

    return {

        "label": label,

        "samples": len(results),

        "chars": total_chars,

        "char_div_4_estimate": total_naive,

        "best_estimate": round(total_best, 1),

        "calibration_factor_vs_char_div_4": round(total_best / total_naive, 3) if total_naive else None,

        "exact_tokenizer_samples": exact_count,

        "method": "exact_tokenizer" if exact_count == len(results) and results else (

            "mixed" if exact_count else "composition_heuristic"

        ),

        "note": (

            "factor > 1.0 means char/4 UNDERESTIMATES real tokens for this output style; "

            "apply the factor before comparing token deltas across subjects with different styles"

        ),

    }

# Representative tool-call framing per provider family. The command itself is

# excluded (zero-length placeholder) so the measurement isolates ENVELOPE

# overhead: what the model pays per call beyond the command text and beyond

# the raw stdout. Used to validate oakbench.tokens envelope constants.

ENVELOPE_TEMPLATES = {

    "anthropic_tool_use": json.dumps(

        {

            "type": "tool_use",

            "id": "toolu_01A1B2C3D4E5F6G7H8J9K0L1",

            "name": "Bash",

            "input": {"command": "", "description": "Run version control command"},

        }

    ),

    "anthropic_tool_result": json.dumps(

        {

            "type": "tool_result",

            "tool_use_id": "toolu_01A1B2C3D4E5F6G7H8J9K0L1",

            "content": "",

        }

    ),

    "openai_function_call": json.dumps(

        {

            "id": "call_a1B2c3D4e5F6g7H8j9K0l1M2",

            "type": "function",

            "function": {"name": "shell", "arguments": json.dumps({"command": [""]})},

        }

    ),

    "openai_function_result": json.dumps(

        {"role": "tool", "tool_call_id": "call_a1B2c3D4e5F6g7H8j9K0l1M2", "content": ""}

    ),

}

def envelope_report() -> dict[str, Any]:

    """Measure tool-call envelope token cost per provider template.

    Compares against the additive envelope constants in oakbench.tokens so a

    drifting provider format is caught by re-running this tool, not by a stale

    constant silently mispricing every row.

    """

    try:

        from oakbench.tokens import (

            TOOL_CALL_ENVELOPE_EMITTED_TOKENS,

            TOOL_RESULT_ENVELOPE_INGESTED_TOKENS,

        )

        constants: dict[str, Any] = {

            "per_call_emitted": TOOL_CALL_ENVELOPE_EMITTED_TOKENS,

            "per_call_ingested": TOOL_RESULT_ENVELOPE_INGESTED_TOKENS,

        }

    except ImportError:

        constants = {}

    samples = {name: analyze_text(name, text) for name, text in ENVELOPE_TEMPLATES.items()}

    return {

        "harness_constants": constants,

        "templates": samples,

        "note": (

            "best_estimate per template is the per-call envelope token cost for that "

            "provider framing, excluding command text and stdout. If these drift far "

            "from harness_constants, update oakbench.tokens and note it in the run metadata."

        ),

    }

def parse_args() -> argparse.Namespace:

    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("paths", nargs="*", type=Path, help="Files or directories of captured command output.")

    parser.add_argument("--label", help="Group label recorded in the aggregate, e.g. subject or operation name.")

    parser.add_argument("--max-bytes", type=int, default=DEFAULT_MAX_BYTES, help="Per-file read cap.")

    parser.add_argument("--per-file", action="store_true", help="Include per-file results, not just the aggregate.")

    parser.add_argument("--self-test", action="store_true", help="Run on embedded diff/hash/prose samples.")

    parser.add_argument(

        "--envelope",

        action="store_true",

        help="Measure tool-call envelope overhead per provider template and compare to harness constants.",

    )

    parser.add_argument("--compact", action="store_true")

    return parser.parse_args()

def main() -> int:

    args = parse_args()

    results: list[dict[str, Any]] = []

    if args.envelope:

        report = envelope_report()

        if args.compact:

            print(json.dumps(report, sort_keys=True, separators=(",", ":")))

        else:

            print(json.dumps(report, indent=2, sort_keys=True))

        return 0

    if args.self_test:

        for name, text in SELF_TEST_SAMPLES.items():

            results.append(analyze_text(name, text))

    elif args.paths:

        for path in iter_input_files(args.paths):

            try:

                text = path.read_bytes()[: args.max_bytes].decode("utf-8", "replace")

            except OSError as exc:

                print(f"skipping {path}: {exc}", file=sys.stderr)

                continue

            if not text.strip():

                continue

            results.append(analyze_text(str(path), text))

    else:

        raise SystemExit("Provide input paths or --self-test. See --help.")

    if not results:

        raise SystemExit("No non-empty input samples found.")

    report: dict[str, Any] = {"aggregate": aggregate(results, args.label)}

    if args.per_file or args.self_test:

        report["samples"] = results

    if args.compact:

        print(json.dumps(report, sort_keys=True, separators=(",", ":")))

    else:

        print(json.dumps(report, indent=2, sort_keys=True))

    return 0

if __name__ == "__main__":

    raise SystemExit(main())