#!/usr/bin/env python3

"""Measure chars-per-token factors per output style and write token_calibration.json.

Rows stay raw (char-based estimates); this campaign produces the versioned

factors file that REPORTING uses to convert raw char counts into calibrated

token estimates via oakbench.calibration. The committed

config/token_calibration.json is an uncalibrated placeholder; this script

overwrites it (or writes elsewhere via --out) with measured factors.

Samples are *.txt files in --samples-dir, categorized by filename prefix

matching a known style, e.g. status_output_01.txt -> status_output. For each

sample the factor is chars / tokens; per style we record the mean factor, a

bootstrap CI (oakbench.stats.bootstrap_ci, statistic "mean"), and n. A style

with no samples keeps null factors with a reason — null means unmeasured

(ADR-0002), never a silent default.

Token-counting methods, in preference order:

- tiktoken o200k_base when importable (optional, never required — ADR-0001)

- Anthropic count-tokens API: requires explicit --method anthropic AND

  ANTHROPIC_API_KEY; not implemented in this slice, and NEVER called

  implicitly — --method auto stays offline.

- composition heuristic: chars/4.0 for prose-like styles, chars/3.2 for

  diff-like styles (BPE-family tokenizers split diff syntax more finely).

Usage:

    python3 scripts/token_calibration_campaign.py \

        --samples-dir results/calibration-samples \

        --method auto \

        --calibration-version cal-2026.06-v1 \

        --out config/token_calibration.json

"""

from __future__ import annotations

import argparse

import json

import statistics

import sys

from pathlib import Path

from typing import Any, Callable

sys.path.insert(0, str(Path(__file__).resolve().parent))

from oakbench.calibration import SCHEMA_VERSION, available_token_counters

from oakbench.stats import bootstrap_ci

STYLES = ("status_output", "diff_output", "log_output", "prose")

# Documented heuristic constants (see module docstring): diff-like text

# tokenizes around 3.2 chars/token, prose-like around 4.0.

HEURISTIC_PROSE_CHARS_PER_TOKEN = 4.0

HEURISTIC_DIFF_CHARS_PER_TOKEN = 3.2

DIFF_LIKE_STYLES = frozenset({"diff_output"})

def heuristic_token_counter(style: str) -> Callable[[str], float]:

    rate = HEURISTIC_DIFF_CHARS_PER_TOKEN if style in DIFF_LIKE_STYLES else HEURISTIC_PROSE_CHARS_PER_TOKEN

    return lambda text: len(text) / rate

def tiktoken_token_counter(style: str) -> Callable[[str], float]:

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

    encoding = tiktoken.get_encoding("o200k_base")

    return lambda text: float(len(encoding.encode(text, disallowed_special=())))

def resolve_method(requested: str) -> tuple[str, Callable[[str], Callable[[str], float]]]:

    """Map the --method flag to a recorded method name and a counter factory.

    Never falls through to a network call: anthropic must be requested

    explicitly, and even then this slice refuses rather than calling out.

    """

    counters = available_token_counters()

    if requested == "anthropic":

        if not counters["anthropic_api"]:

            raise SystemExit("--method anthropic requires ANTHROPIC_API_KEY in the environment")

        raise SystemExit("--method anthropic: not implemented in this slice")

    if requested == "tiktoken":

        if not counters["tiktoken"]:

            raise SystemExit("--method tiktoken requires the optional tiktoken package (not importable)")

        return "tiktoken_o200k_base", tiktoken_token_counter

    if requested == "heuristic":

        return "composition_heuristic", heuristic_token_counter

    if requested == "auto":

        if counters["tiktoken"]:

            return "tiktoken_o200k_base", tiktoken_token_counter

        return "composition_heuristic", heuristic_token_counter

    raise SystemExit(f"unknown --method {requested!r}")

def collect_samples(samples_dir: Path) -> dict[str, list[Path]]:

    """Group *.txt samples by known-style filename prefix; warn on strays."""

    if not samples_dir.is_dir():

        raise SystemExit(f"--samples-dir not found or not a directory: {samples_dir}")

    by_style: dict[str, list[Path]] = {style: [] for style in STYLES}

    for path in sorted(samples_dir.glob("*.txt")):

        stem = path.stem

        for style in STYLES:

            if stem == style or stem.startswith(style + "_"):

                by_style[style].append(path)

                break

        else:

            print(f"skipping {path.name}: no known style prefix ({', '.join(STYLES)})", file=sys.stderr)

    return by_style

def measure_style(

    paths: list[Path], style: str, counter_factory: Callable[[str], Callable[[str], float]]

) -> dict[str, Any]:

    """Per-style factor entry: mean chars/tokens, bootstrap CI, n.

    Empty or zero-token samples are skipped with a warning; a style that ends

    up with no usable samples keeps null factors with a reason (ADR-0002).

    """

    count_tokens = counter_factory(style)

    factors: list[float] = []

    for path in paths:

        text = path.read_text(encoding="utf-8", errors="replace")

        chars = len(text)

        if chars == 0:

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

            continue

        tokens = count_tokens(text)

        if tokens <= 0:

            print(f"skipping {path.name}: zero tokens counted", file=sys.stderr)

            continue

        factors.append(chars / tokens)

    if not factors:

        return {

            "chars_per_token": None,

            "ci_low": None,

            "ci_high": None,

            "n": 0,

            "reason": "no samples" if not paths else "no usable samples",

        }

    ci_low, ci_high = bootstrap_ci(factors, statistic="mean")

    return {

        "chars_per_token": float(statistics.mean(factors)),

        "ci_low": float(ci_low),

        "ci_high": float(ci_high),

        "n": len(factors),

    }

def parse_args() -> argparse.Namespace:

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

    parser.add_argument("--samples-dir", type=Path, required=True, help="Directory of style-prefixed *.txt samples.")

    parser.add_argument(

        "--method",

        choices=("auto", "tiktoken", "heuristic", "anthropic"),

        default="auto",

        help="Token counter; auto prefers tiktoken when importable and never calls network APIs.",

    )

    parser.add_argument("--out", type=Path, required=True, help="Path for the new token_calibration.json.")

    parser.add_argument(

        "--calibration-version",

        required=True,

        help="Version string stamped into the file and cited by every calibrated number, e.g. cal-2026.06-v1.",

    )

    return parser.parse_args()

def main() -> int:

    args = parse_args()

    method_name, counter_factory = resolve_method(args.method)

    by_style = collect_samples(args.samples_dir)

    factors = {style: measure_style(paths, style, counter_factory) for style, paths in by_style.items()}

    measured_ns = [entry["n"] for entry in factors.values() if entry["n"] > 0]

    document = {

        "schema_version": SCHEMA_VERSION,

        "calibration_version": args.calibration_version,

        "method": method_name,

        "samples_per_style": min(measured_ns) if measured_ns else 0,

        "factors": factors,

    }

    args.out.parent.mkdir(parents=True, exist_ok=True)

    args.out.write_text(json.dumps(document, indent=2, sort_keys=True) + "\n")

    if not measured_ns:

        print(f"wrote {args.out}, but no styles were measured", file=sys.stderr)

        return 1

    unmeasured = sorted(style for style, entry in factors.items() if entry["n"] == 0)

    print(

        f"wrote {args.out}: {len(measured_ns)}/{len(STYLES)} styles measured via {method_name}"

        + (f" (unmeasured: {', '.join(unmeasured)})" if unmeasured else "")

    )

    return 0

if __name__ == "__main__":

    raise SystemExit(main())