#!/usr/bin/env python3
"""
Measure how many LLM input tokens a Dragon-mode message costs in Simplified Chinese
(zh-Hans) compared with the SAME message in the user's own language.

Input : scripts/token-corpus.json   (12 parallel messages per language + zh-Hans)
Output: src/site/token-data.json    (imported by the landing-page token calculator)

Tokenizers:
  * o200k_base  - OpenAI GPT-4o / GPT-4.1 / GPT-5 family (open, via tiktoken)
  * cl100k_base - OpenAI GPT-4 / GPT-3.5-turbo era      (open, via tiktoken)
  * Anthropic   - ONLY if ANTHROPIC_API_KEY and ANTHROPIC_COUNT_MODEL are set; uses the
                  free POST /v1/messages/count_tokens endpoint. The per-request overhead is
                  removed by subtracting the count for a one-character baseline message.

Usage (from apps/web):
  python -m pip install tiktoken
  python scripts/measure-tokens.py
  # standalone (files downloaded from https://dragonwispr.com/methodology):
  python measure-tokens.py --corpus token-corpus.json --out token-data.json
  # optional Claude numbers:
  ANTHROPIC_API_KEY=... ANTHROPIC_COUNT_MODEL=<claude model id> python scripts/measure-tokens.py

MARKETING GUARDRAIL (binding, see CLAUDE.md): claims compare zh-Hans against the user's OWN
language only. For English the data shows Chinese costs MORE tokens; the site must say so.
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import urllib.request
from datetime import date
from pathlib import Path

import tiktoken

HERE = Path(__file__).resolve().parent
CORPUS = HERE / "token-corpus.json"
OUT = HERE.parent / "src" / "site" / "token-data.json"

TARGET = "zh-Hans"
SOURCE_ORDER = ["el", "de", "fr", "es", "it", "pt", "ru", "uk", "pl", "tr",
                "ar", "he", "hi", "ja", "ko", "vi", "en"]
# Languages written without spaces between words: a whitespace word count is meaningless.
NO_SPACE_SCRIPTS = {"ja", "zh-Hans"}
EXPECTED_MESSAGES = 12
TOKENIZERS = {
    "o200k_base": "OpenAI o200k_base (GPT-4o, GPT-4.1, GPT-5 family)",
    "cl100k_base": "OpenAI cl100k_base (GPT-4, GPT-3.5-turbo)",
}


def anthropic_counter():
    key = os.environ.get("ANTHROPIC_API_KEY")
    model = os.environ.get("ANTHROPIC_COUNT_MODEL")
    if not key or not model:
        return None, None

    def raw(text: str) -> int:
        body = json.dumps({"model": model, "messages": [{"role": "user", "content": text}]})
        req = urllib.request.Request(
            "https://api.anthropic.com/v1/messages/count_tokens",
            data=body.encode("utf-8"),
            headers={
                "x-api-key": key,
                "anthropic-version": "2023-06-01",
                "content-type": "application/json",
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=30) as resp:
            return int(json.load(resp)["input_tokens"])

    baseline = raw(".") - 1  # request/turn overhead; "." itself is assumed to be 1 token

    def count(text: str) -> int:
        return raw(text) - baseline

    return count, model


def pct(new: int, old: int) -> float:
    return round((new - old) / old * 100, 1)


def main() -> int:
    parser = argparse.ArgumentParser(description="Measure zh-Hans token counts against each source language.")
    parser.add_argument("--corpus", type=Path, default=CORPUS, help=f"parallel corpus (default: {CORPUS})")
    parser.add_argument("--out", type=Path, default=OUT, help=f"output JSON (default: {OUT})")
    args = parser.parse_args()
    out_path: Path = args.out

    corpus = json.loads(args.corpus.read_text(encoding="utf-8"))["languages"]
    for code in SOURCE_ORDER + [TARGET]:
        n = len(corpus[code]["messages"])
        if n != EXPECTED_MESSAGES:
            sys.exit(f"{code}: expected {EXPECTED_MESSAGES} messages, got {n}")
    en_words = [len(m.split()) for m in corpus["en"]["messages"]]
    if not all(15 <= w <= 60 for w in en_words):
        sys.exit(f"English master messages must be 15-60 words, got {en_words}")

    encoders = {name: tiktoken.get_encoding(name) for name in TOKENIZERS}
    counters = {name: (lambda e: (lambda t: len(e.encode(t))))(enc) for name, enc in encoders.items()}
    claude_count, claude_model = anthropic_counter()
    if claude_count:
        counters["anthropic"] = claude_count

    def per_message(code: str, tok: str) -> list[int]:
        return [counters[tok](m) for m in corpus[code]["messages"]]

    zh_counts = {tok: per_message(TARGET, tok) for tok in counters}
    zh_msgs = corpus[TARGET]["messages"]

    languages = []
    for code in SOURCE_ORDER:
        msgs = corpus[code]["messages"]
        entry = {
            "code": code,
            "name": corpus[code]["name"],
            "nativeName": corpus[code]["nativeName"],
            "messages": len(msgs),
            "chars": sum(len(m) for m in msgs),
            "zhChars": sum(len(m) for m in zh_msgs),
            "wordsPerMessage": None if code in NO_SPACE_SCRIPTS
            else round(sum(len(m.split()) for m in msgs) / len(msgs), 1),
            "tokens": {},
        }
        for tok in counters:
            src = per_message(code, tok)
            zh = zh_counts[tok]
            per_msg_pct = [pct(z, s) for z, s in zip(zh, src)]
            entry["tokens"][tok] = {
                "source": sum(src),
                "zh": sum(zh),
                "ratioZhToSource": round(sum(zh) / sum(src), 3),
                "pctChange": pct(sum(zh), sum(src)),
                "perMessagePctChange": {
                    "min": min(per_msg_pct),
                    "median": round(statistics.median(per_msg_pct), 1),
                    "max": max(per_msg_pct),
                },
            }
        primary = entry["tokens"]["o200k_base"]["pctChange"]
        entry["chineseSavesTokens"] = all(t["pctChange"] < 0 for t in entry["tokens"].values())
        entry["verdict"] = (
            f"Chinese uses {abs(primary):.0f}% fewer tokens than {entry['name']} (o200k_base)"
            if primary < 0
            else f"Chinese uses {primary:.0f}% MORE tokens than {entry['name']} (o200k_base) - no saving"
        )
        languages.append(entry)

    data = {
        "version": 1,
        "measuredOn": date.today().isoformat(),
        "methodology": {
            "summary": (
                "12 realistic messages to an AI assistant / coding agent were written in English, "
                "then translated faithfully into each language and into Simplified Chinese (zh-Hans) "
                "as Dragon mode would output it. Every message was tokenized separately and the totals "
                "were compared: pctChange = (zh - source) / source * 100. Negative means Chinese uses fewer tokens."
            ),
            "corpus": {
                "file": "apps/web/scripts/token-corpus.json",
                "messagesPerLanguage": EXPECTED_MESSAGES,
                "englishWordsPerMessage": {"min": min(en_words), "max": max(en_words)},
                "domainMix": "4 coding requests with English identifiers, 2 writing help, 2 analysis, "
                             "2 everyday, 1 CI/devops, 1 personal finance explanation",
                "conventions": "Code identifiers, file paths and product names stay in Latin script in every "
                               "language (as Dragon mode will keep them). Each language uses its own natural "
                               "conventions for numerals and common technical loanwords.",
            },
            "script": "apps/web/scripts/measure-tokens.py",
            "tokenizers": {
                **TOKENIZERS,
                "library": f"tiktoken {tiktoken.__version__}",
                **({"anthropic": f"Anthropic count_tokens API, model {claude_model}"} if claude_model else {}),
            },
            "claudeIncluded": bool(claude_count),
            "claudeNote": None if claude_count else (
                "Claude (Anthropic) token counts are NOT included: no ANTHROPIC_API_KEY was available "
                "when this file was generated. Re-run the script with ANTHROPIC_API_KEY and "
                "ANTHROPIC_COUNT_MODEL set to add them."
            ),
            "caveats": [
                "Small corpus (12 messages per language); treat results as indicative, not as a benchmark.",
                "Translations were written by hand for this measurement; a different translator or Dragon mode's "
                "live model output may be longer or shorter, especially for zh-Hans.",
                "Results depend on the domain mix: messages full of code identifiers save less, because "
                "identifiers stay in Latin script and cost the same in every language.",
                "Only OpenAI's open tokenizers are measured; other models (Claude, Gemini, Llama, Qwen) use "
                "different tokenizers and will give different numbers.",
                "Only input tokens are measured. Model replies are not part of this comparison.",
                "For English, Chinese costs MORE tokens; Dragon mode is not a token saver for English speakers.",
            ],
            "guardrail": "Every public claim compares zh-Hans against the user's OWN language, never against English.",
        },
        "target": {"code": TARGET, "name": corpus[TARGET]["name"]},
        "languages": languages,
    }

    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    print(f"wrote {out_path}")
    hdr = f"{'lang':6}{'chars':>7}{'w/msg':>7}"
    for tok in counters:
        hdr += f" | {tok:>11} src  zh   ratio   pct"
    print(hdr)
    for e in languages:
        row = f"{e['code']:6}{e['chars']:>7}{(e['wordsPerMessage'] or 0):>7}"
        for tok in counters:
            t = e["tokens"][tok]
            row += f" | {'':>11}{t['source']:>4} {t['zh']:>4} {t['ratioZhToSource']:>6} {t['pctChange']:>6}"
        print(row)
    return 0


if __name__ == "__main__":
    sys.exit(main())