"""
Count the same 50 text samples with two OpenAI tokenizers (tiktoken) and with
each of the four Gemini model ids that appear on getllmcalc.com, via the
Gemini API countTokens endpoint (a free endpoint, no generation involved).

The samples are NOT duplicated here: they are imported from
measure_chars_per_token.py, which must sit in the same directory. Download
both files together to re-run. Samples are self-authored (5 categories x 10
segments: english_prose, chinese_traditional, chinese_simplified,
python_code, json).

Rules this script follows:
  * Each Gemini model id is called on its own. A model that fails is recorded
    as unavailable (status "n/a"); there is no fallback to any other model,
    and one model's ratio is never reused for another model.
  * GPT side = tiktoken o200k_base and cl100k_base. Applying o200k_base to
    the GPT-5 family is an assumption (the site's token counter makes the
    same assumption); it is not a measurement of those models.
  * Claude is not measured (no API key available for its counter).

Gemini overhead: countTokens adds a fixed overhead per request (an empty
string is counted as 1 token). The script measures that baseline per model
with an empty string and reports counts both raw and net (raw - baseline).
All ratios and costs in the article use the net counts.

Aggregation (same as measure_chars_per_token.py): per-category average of the
per-sample chars/token ratios, rounded to 3 decimals. Cost figures use the
pooled ratio (total chars / total tokens) because cost depends on totals.

Run:
    pip install tiktoken google-genai
    set GEMINI_API_KEY=...        (your own key; never put it in a file)
    python measure_gemini_vs_gpt_tokens.py

Writes: measure_gemini_vs_gpt_tokens.results.json (next to this script).
"""

import json
import os
import statistics
import sys
import time
from datetime import date
from pathlib import Path

import tiktoken

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from measure_chars_per_token import CATEGORIES  # noqa: E402  (same 50 samples)

PRICING_PATH = HERE.parent / "pricing.json"
OUT_PATH = HERE / "measure_gemini_vs_gpt_tokens.results.json"

GPT_ENCODINGS = ["cl100k_base", "o200k_base"]
# The four Gemini ids listed on the site (pricing.json ids minus the "gemini/" prefix).
GEMINI_MODELS = [
    "gemini-3.1-pro-preview",
    "gemini-3.6-flash",
    "gemini-3.5-flash-lite",
    "gemini-3.1-flash-lite",
]
RETRIES = 4
PAUSE_S = 0.15


def count_gemini(client, model, text):
    """Return total_tokens for `text` on `model`, retrying transient errors."""
    last = None
    for attempt in range(RETRIES):
        try:
            resp = client.models.count_tokens(model=model, contents=text)
            return int(resp.total_tokens)
        except Exception as exc:  # network / quota / model errors
            last = exc
            time.sleep(1.5 * (attempt + 1))
    raise RuntimeError(type(last).__name__ + ": " + str(last)[:200])


def measure_gemini_model(client, model):
    """Count all 50 samples + the empty-string baseline for one model id."""
    out = {"requested_model_id": model}
    try:
        baseline = count_gemini(client, model, "")
        raw = {}
        for category, samples in CATEGORIES.items():
            counts = []
            for text in samples:
                counts.append(count_gemini(client, model, text))
                time.sleep(PAUSE_S)
            raw[category] = counts
    except Exception as exc:
        out["status"] = "n/a"
        out["error"] = str(exc).encode("ascii", "replace").decode()
        return out
    out["status"] = "ok"
    out["baseline_tokens_empty_string"] = baseline
    out["raw_counts"] = raw
    return out


def ratio(chars, tokens):
    return chars / tokens if tokens else 0


def main():
    key = os.environ.get("GEMINI_API_KEY")
    if not key:
        print("GEMINI_API_KEY is not set; cannot call countTokens.")
        sys.exit(1)
    from google import genai
    client = genai.Client(api_key=key)

    encoders = {n: tiktoken.get_encoding(n) for n in GPT_ENCODINGS}
    gemini_runs = {m: measure_gemini_model(client, m) for m in GEMINI_MODELS}
    ok_models = [m for m in GEMINI_MODELS if gemini_runs[m]["status"] == "ok"]

    pricing = json.loads(PRICING_PATH.read_text(encoding="utf-8"))
    rates = {row["id"]: row for row in pricing["chat"]}

    results = {
        "measured_on": date.today().isoformat(),
        "gpt_encodings_measured": GPT_ENCODINGS,
        "gemini_models_requested": GEMINI_MODELS,
        "gemini_models_measured": ok_models,
        "gemini_model_status": {
            m: {k: v for k, v in gemini_runs[m].items() if k != "raw_counts"}
            for m in GEMINI_MODELS
        },
        "gemini_response_model_id_note": (
            "countTokens returns only total_tokens; the response carries no "
            "model id, so the requested id is the only id recorded."
        ),
        "claude_measured": False,
        "categories": {},
    }

    for category, samples in CATEGORIES.items():
        chars = [len(t) for t in samples]
        per_tok = {}  # tokenizer label -> list of token counts (net for Gemini)
        for n in GPT_ENCODINGS:
            per_tok[n] = [len(encoders[n].encode(t)) for t in samples]
        for m in ok_models:
            base = gemini_runs[m]["baseline_tokens_empty_string"]
            per_tok[m] = [c - base for c in gemini_runs[m]["raw_counts"][category]]

        cat = {"sample_count": len(samples), "total_chars": sum(chars), "tokenizers": {}}
        for label, toks in per_tok.items():
            ratios = [ratio(c, t) for c, t in zip(chars, toks)]
            entry = {
                "avg_chars_per_token": round(statistics.mean(ratios), 3),
                "pooled_chars_per_token": round(ratio(sum(chars), sum(toks)), 3),
                "total_tokens": sum(toks),
                "tokens_vs_o200k": round(sum(toks) / sum(per_tok["o200k_base"]), 3),
            }
            if label in gemini_runs:
                raw = gemini_runs[label]["raw_counts"][category]
                entry["total_tokens_raw_with_overhead"] = sum(raw)
            entry["samples"] = [
                {"chars": c, "tokens": t, "chars_per_token": round(ratio(c, t), 3)}
                for c, t in zip(chars, toks)
            ]
            cat["tokenizers"][label] = entry
        results["categories"][category] = cat

    # Input cost per 1M characters = in_rate * total_tokens / total_chars (USD), i.e.
# in_rate / pooled chars per token computed from the exact totals.
    # GPT rows: o200k_base assumed (not measured on the GPT-5 models themselves).
    # Gemini rows: that model's own ratio only. Claude: not estimated.
    cost = []
    for row in pricing["chat"]:
        entry = {"id": row["id"], "name": row["name"], "vendor": row["vendor"],
                 "in_per_million_tokens": row["in"], "basis": None, "usd_per_million_chars": {}}
        if row["vendor"] == "OpenAI":
            entry["basis"] = "o200k_base (assumed for this model)"
            label = "o200k_base"
        elif row["vendor"] == "Google" and row["id"].startswith("gemini/") \
                and row["id"][len("gemini/"):] in ok_models:
            label = row["id"][len("gemini/"):]
            entry["basis"] = label + " (measured on this model id)"
        else:
            entry["basis"] = "n/a"
            label = None
        for category in CATEGORIES:
            if label is None:
                entry["usd_per_million_chars"][category] = None
            else:
                t = results["categories"][category]["tokenizers"][label]
                # exact totals, not the rounded pooled ratio (avoids double rounding)
                entry["usd_per_million_chars"][category] = round(
                    row["in"] * t["total_tokens"] / results["categories"][category]["total_chars"], 4)
        cost.append(entry)
    results["pricing_last_synced"] = pricing.get("last_synced")
    results["input_cost_per_million_chars"] = cost

    # Pairs with the same listed input rate (USD per 1M tokens): the cost ratio
    # Gemini / GPT per character is then just Gemini total tokens / o200k_base
    # total tokens (tokens_vs_o200k), so it shows the tokenizer effect alone.
    pairs = []
    for g in cost:
        if g["vendor"] != "Google" or g["basis"] == "n/a":
            continue
        for o in cost:
            if o["vendor"] == "OpenAI" and o["in_per_million_tokens"] == g["in_per_million_tokens"]:
                pairs.append({
                    "gemini": g["name"], "gpt": o["name"],
                    "shared_in_per_million_tokens": g["in_per_million_tokens"],
                    "gemini_cost_over_gpt_cost": {
                        c: results["categories"][c]["tokenizers"][g["id"][len("gemini/"):]]["tokens_vs_o200k"]
                        for c in CATEGORIES
                    },
                })
    results["same_input_rate_pairs"] = pairs

    OUT_PATH.write_text(json.dumps(results, ensure_ascii=False, indent=2) + "\n",
                        encoding="utf-8", newline="\n")

    print("Wrote", OUT_PATH.name)
    print("gemini models measured:", ", ".join(ok_models) or "none")
    for m in GEMINI_MODELS:
        s = gemini_runs[m]
        print(" ", m, s["status"], s.get("baseline_tokens_empty_string", s.get("error", "")))
    print()
    for category, cat in results["categories"].items():
        for label, e in cat["tokenizers"].items():
            print("%-20s %-24s avg=%7.3f pooled=%7.3f tokens=%5d vs_o200k=%.3f"
                  % (category, label, e["avg_chars_per_token"], e["pooled_chars_per_token"],
                     e["total_tokens"], e["tokens_vs_o200k"]))


if __name__ == "__main__":
    main()