#!/usr/bin/env python3
"""Токен-калькулятор TextMachine: замер фактической токенизации художественных текстов.

Считает для каждого файла eval/data/samples/<lang>/*.txt метрики по всем токенизаторам:
  - tokens/CJK-знак (zh, ja), tokens/слово (en, ru), символов/токен;
  - при наличии manifest.json с parallel_of — соотношение объёмов перевод/оригинал.

Запуск:  eval/.venv/bin/python eval/token_calc.py [--samples-dir eval/data/samples] [--json-out ...]
Выход: markdown-таблицы в stdout + JSON с сырыми числами.
"""
from __future__ import annotations

import argparse
import json
import re
import sys
import unicodedata
from dataclasses import dataclass, field
from pathlib import Path

ROOT = Path(__file__).resolve().parent

# --- Токенизаторы -----------------------------------------------------------

def load_tokenizers() -> dict[str, "callable"]:
    """Возвращает {имя: fn(text)->int}. Каждый fn считает токены."""
    toks: dict[str, callable] = {}

    import tiktoken
    for enc_name, label in [("o200k_base", "gpt-o200k (GPT-4o/5)"),
                            ("cl100k_base", "gpt-cl100k (GPT-4)")]:
        enc = tiktoken.get_encoding(enc_name)
        toks[label] = (lambda e: lambda t: len(e.encode(t, disallowed_special=())))(enc)

    from tokenizers import Tokenizer
    hf_dirs = {
        "deepseek-v3.2": "DeepSeek V3.2",
        "qwen3": "Qwen3",
        "qwen2.5": "Qwen2.5",
        "glm-4.6": "GLM-4.6",
    }
    for d, label in hf_dirs.items():
        path = ROOT / "tokenizers" / d / "tokenizer.json"
        if path.exists():
            tk = Tokenizer.from_file(str(path))
            toks[label] = (lambda t_: lambda t: len(t_.encode(t, add_special_tokens=False).ids))(tk)
        else:
            print(f"[warn] нет {path}, пропускаю {label}", file=sys.stderr)
    return toks


# --- Подсчёт символов/слов ---------------------------------------------------

CJK_RE = re.compile(
    "[㐀-䶿一-鿿豈-﫿\U00020000-\U0002ebef]"
)  # хань-иероглифы (кандзи/ханьцзы), вкл. расширения
KANA_RE = re.compile("[぀-ヿ]")  # хирагана + катакана
WORD_RE = re.compile(r"[A-Za-zА-Яа-яЁёÀ-ɏ]+(?:[-'’][A-Za-zА-Яа-яЁё]+)*")


@dataclass
class TextStats:
    path: str
    lang: str
    chars_total: int = 0
    chars_no_ws: int = 0
    cjk_chars: int = 0      # хань-иероглифы
    kana_chars: int = 0     # ja: кана
    words: int = 0          # en/ru: слова
    tokens: dict = field(default_factory=dict)  # {tokenizer: n}


def analyze(path: Path, lang: str, toks: dict) -> TextStats:
    text = path.read_text(encoding="utf-8")
    # нормализуем переводы строк, убираем BOM
    text = text.replace("", "").strip()
    st = TextStats(path=str(path), lang=lang)
    st.chars_total = len(text)
    st.chars_no_ws = sum(1 for c in text if not c.isspace())
    st.cjk_chars = len(CJK_RE.findall(text))
    st.kana_chars = len(KANA_RE.findall(text))
    st.words = len(WORD_RE.findall(text))
    for name, fn in toks.items():
        st.tokens[name] = fn(text)
    return st


def unit_count(st: TextStats) -> tuple[str, int]:
    """Опорная единица языка: zh/ja — CJK-знаки (для ja — кандзи+кана), en/ru — слова."""
    if st.lang == "zh":
        return "иероглиф", st.cjk_chars
    if st.lang == "ja":
        return "знак (кандзи+кана)", st.cjk_chars + st.kana_chars
    return "слово", st.words


# --- Отчёт -------------------------------------------------------------------

def fmt_table(rows: list[list[str]], header: list[str]) -> str:
    out = ["| " + " | ".join(header) + " |",
           "|" + "|".join("---" for _ in header) + "|"]
    out += ["| " + " | ".join(r) + " |" for r in rows]
    return "\n".join(out)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--samples-dir", default=str(ROOT / "data" / "samples"))
    ap.add_argument("--json-out", default=str(ROOT / "data" / "token_calc_results.json"))
    args = ap.parse_args()

    samples = Path(args.samples_dir)
    toks = load_tokenizers()
    tok_names = list(toks)

    stats: list[TextStats] = []
    for lang_dir in sorted(samples.iterdir()):
        if not lang_dir.is_dir():
            continue
        lang = lang_dir.name
        for f in sorted(lang_dir.glob("*.txt")):
            stats.append(analyze(f, lang, toks))

    if not stats:
        sys.exit(f"Нет файлов в {samples}/<lang>/*.txt")

    # 1. Таблица per-file: tokens/единица по каждому токенизатору
    print("## Токены на опорную единицу языка (zh/ja — знак, en/ru — слово)\n")
    header = ["Файл", "Язык", "Единиц"] + tok_names
    rows = []
    for st in stats:
        unit, n = unit_count(st)
        rows.append([Path(st.path).name, st.lang, f"{n}"]
                    + [f"{st.tokens[t] / n:.3f}" if n else "—" for t in tok_names])
    print(fmt_table(rows, header))

    # 2. Агрегат по языку (взвешенно по объёму)
    print("\n## Агрегат по языку (сумма токенов / сумма единиц)\n")
    rows = []
    langs = sorted({st.lang for st in stats})
    agg: dict[str, dict] = {}
    for lang in langs:
        group = [st for st in stats if st.lang == lang]
        n_units = sum(unit_count(st)[1] for st in group)
        n_chars = sum(st.chars_no_ws for st in group)
        row = [lang, unit_count(group[0])[0], f"{n_units}"]
        agg[lang] = {"unit": unit_count(group[0])[0], "units": n_units,
                     "chars_no_ws": n_chars, "per_unit": {}, "chars_per_token": {}}
        for t in tok_names:
            total = sum(st.tokens[t] for st in group)
            row.append(f"{total / n_units:.3f}")
            agg[lang]["per_unit"][t] = total / n_units
            agg[lang]["chars_per_token"][t] = n_chars / total
        rows.append(row)
    print(fmt_table(rows, ["Язык", "Единица", "Единиц"] + tok_names))

    print("\n## Символов (без пробелов) на токен\n")
    rows = [[lang, *(f"{agg[lang]['chars_per_token'][t]:.2f}" for t in tok_names)]
            for lang in langs]
    print(fmt_table(rows, ["Язык"] + tok_names))

    # 3. Параллельные пары из manifest.json
    manifest_path = samples / "manifest.json"
    pairs_out = []
    if manifest_path.exists():
        manifest = json.loads(manifest_path.read_text())
        by_path = {st.path: st for st in stats}
        entries = manifest if isinstance(manifest, list) else manifest.get("files", [])
        pairs = [(e["path"], e["parallel_of"]) for e in entries if e.get("parallel_of")]
        if pairs:
            print("\n## Параллельные пары: объём перевода к оригиналу (в токенах одного токенизатора)\n")
            rows = []

            def resolve(p):
                return by_path.get(str(ROOT.parent / p)) or by_path.get(p)

            for tr_path, src_spec in pairs:
                # parallel_of может быть списком файлов (перевод покрывает несколько)
                src_paths = src_spec if isinstance(src_spec, list) else [src_spec]
                tr = resolve(tr_path)
                srcs = [resolve(s) for s in src_paths]
                if not tr or not all(srcs):
                    print(f"[warn] пара не найдена: {tr_path} / {src_spec}", file=sys.stderr)
                    continue
                src_tokens = {t: sum(s.tokens[t] for s in srcs) for t in tok_names}
                ratios = {t: tr.tokens[t] / src_tokens[t] for t in tok_names}
                pairs_out.append({"translation": tr.path, "source": src_paths,
                                  "src_lang": srcs[0].lang, "ratios": ratios})
                rows.append([Path(tr.path).name,
                             "+".join(Path(s).name for s in src_paths)]
                            + [f"{ratios[t]:.2f}" for t in tok_names])
            if rows:
                print(fmt_table(rows, ["Перевод", "Оригинал"] + tok_names))

    out = {
        "files": [vars(st) for st in stats],
        "aggregate_by_lang": agg,
        "parallel_pairs": pairs_out,
        "tokenizers": tok_names,
    }
    Path(args.json_out).write_text(json.dumps(out, ensure_ascii=False, indent=1))
    print(f"\nJSON: {args.json_out}", file=sys.stderr)


if __name__ == "__main__":
    main()