from pathlib import Path
import math
import os
import random

import pandas as pd
import tiktoken
from datasets import get_dataset_config_names, load_dataset
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer as CoreTokenizer
from transformers import AutoTokenizer


# ---------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------

REPO_ROOT = Path(__file__).resolve().parents[1]

MY_TOKENIZER_PATH = Path(
    os.environ.get("SAMAI_TOKENIZER_PATH", REPO_ROOT / "tokenizer.json")
)

# Four complementary Arabic benchmarks. Each loader uses evaluation/test data
# where available, and no prompt/chat-template text is injected.
EXPECTED_BENCHMARKS = {
    "ArabicMMLU",
    "MLQA-Arabic",
    "MASSIVE-Arabic",
    "Arabic-Dialects",
}

# None = use every available example. Set (for example) 1000 for a quick run;
# the limit is applied independently to each benchmark.
MAX_EXAMPLES_PER_BENCHMARK = None

BATCH_SIZE = 128
RANDOM_SEED = 42

OUTPUT_DIR = Path(
    os.environ.get(
        "SAMAI_BENCHMARK_OUTPUT",
        REPO_ROOT / "results" / "generated",
    )
)

# Some third-party tokenizers require executing code from their repositories.
# Keep this disabled by default. Review those repositories and enable it only
# when you explicitly accept that risk:
#   SAMAI_TRUST_REMOTE_CODE=1 python benchmark/benchmark.py
TRUST_REMOTE_CODE = os.environ.get(
    "SAMAI_TRUST_REMOTE_CODE", "0"
).lower() in {"1", "true", "yes"}

# Only tokenizer/configuration files are downloaded, not model weights.
#
# The list deliberately keeps older baselines while adding newer model
# families. Some repositories (for example Jais, Gemma, Llama and Aya) may
# require accepting their Hugging Face terms and running `hf auth login`.
# A failed or gated tokenizer is reported and skipped without stopping the
# rest of the benchmark.
HF_TOKENIZERS = {
    # Existing Arabic and multilingual baselines
    "Fanar-1": "QCRI/Fanar-1-9B-Instruct",
    "ALLaM-7B": "humain-ai/ALLaM-7B-Instruct-preview",
    "Jais-13B": "inception42/jais-13b",
    "Qwen2.5-7B": "Qwen/Qwen2.5-7B",

    # Newer Qwen tokenizer generations
    "Qwen3-8B": "Qwen/Qwen3-8B",
    "Qwen3.5-4B": "Qwen/Qwen3.5-4B",
    "Qwen3.8-Flash-Next": "Qwen/Qwen3.8-Flash-Next",

    # Recent frontier/open-weight model families
    "DeepSeek-V3": "deepseek-ai/DeepSeek-V3",
    "GPT-OSS-20B": "openai/gpt-oss-20b",
    "GLM-5.3-Flash": "zai-org/GLM-5.3-Flash",
    "Kimi-K2.6": "moonshotai/Kimi-K2.6",

    # Recent multilingual tokenizers with Arabic coverage/relevance
    "Gemma-3-4B": "google/gemma-3-4b-it",
    "Llama-4-Scout": "meta-llama/Llama-4-Scout-17B-16E-Instruct",
    "Mistral-Small-3.1": "mistralai/Mistral-Small-3.1-24B-Instruct-2503",
    "Falcon-H1-7B": "tiiuae/Falcon-H1-7B-Instruct",
    "Aya-Expanse-8B": "CohereLabs/aya-expanse-8b",

    # Additional Arabic-focused baseline
    "AceGPT-v2-8B": "FreedomIntelligence/AceGPT-v2-8B",
}

TIKTOKEN_ENCODINGS = {
    "tiktoken-cl100k": "cl100k_base",
    "tiktoken-o200k": "o200k_base",
}


# ---------------------------------------------------------------------
# Dataset preparation
# ---------------------------------------------------------------------

def clean_value(value) -> str:
    """Convert dataset values to clean text while ignoring null/NaN values."""
    if value is None:
        return ""

    if isinstance(value, float) and math.isnan(value):
        return ""

    value = str(value).strip()

    if value.lower() in {"nan", "none", "null"}:
        return ""

    return value


def arabic_mmlu_row_to_text(row) -> str:
    """
    Build one benchmark document from the original dataset text.

    No artificial English prompt or chat template is added because we want
    to measure the tokenizers, not differences in prompt formatting.
    """
    parts = []

    context = clean_value(row.get("Context"))
    question = clean_value(row.get("Question"))

    if context:
        parts.append(context)

    if question:
        parts.append(question)

    for index in range(1, 6):
        option = clean_value(row.get(f"Option {index}"))
        if option:
            parts.append(option)

    return "\n".join(parts)


def make_record(benchmark, category, text, **metadata):
    text = clean_value(text)
    if not text:
        return None

    return {
        "benchmark": benchmark,
        "category": category,
        "configuration": metadata.get("configuration", ""),
        "group": metadata.get("group", category),
        "subject": metadata.get("subject", category),
        "country": metadata.get("country", ""),
        "text": text,
        "words": len(text.split()),
        "characters": len(text),
        "bytes": len(text.encode("utf-8")),
    }


def deduplicate_and_limit(records):
    unique = []
    seen = set()

    for record in records:
        if record is None or record["text"] in seen:
            continue
        seen.add(record["text"])
        unique.append(record)

    if (
        MAX_EXAMPLES_PER_BENCHMARK is not None
        and len(unique) > MAX_EXAMPLES_PER_BENCHMARK
    ):
        random.Random(RANDOM_SEED).shuffle(unique)
        unique = unique[:MAX_EXAMPLES_PER_BENCHMARK]

    return unique


def load_arabic_mmlu():
    dataset_id = "MBZUAI/ArabicMMLU"
    print(f"\nLoading ArabicMMLU: {dataset_id}")

    configurations = get_dataset_config_names(dataset_id)

    # Some dataset versions may provide an aggregate "All" configuration.
    all_config = next(
        (name for name in configurations if name.lower() == "all"),
        None,
    )

    configs_to_load = [all_config] if all_config else configurations
    records = []

    for position, config_name in enumerate(configs_to_load, start=1):
        print(
            f"[{position}/{len(configs_to_load)}] "
            f"Loading configuration: {config_name}"
        )

        dataset = load_dataset(
            dataset_id,
            config_name,
            split="test",
        )

        for row in dataset:
            # Keep held-out test questions only.
            if row.get("is_few_shot", 0) not in {0, False, None}:
                continue

            records.append(
                make_record(
                    "ArabicMMLU",
                    "academic-exams",
                    arabic_mmlu_row_to_text(row),
                    configuration=config_name,
                    group=clean_value(row.get("Group")) or "Unknown",
                    subject=clean_value(row.get("Subject")) or config_name,
                    country=clean_value(row.get("Country")) or "Unknown",
                )
            )

    return deduplicate_and_limit(records)


def load_mlqa_arabic():
    """Arabic questions and long answer-bearing passages from MLQA."""
    dataset_id = "mteb/MLQARetrieval"
    print(f"\nLoading MLQA Arabic: {dataset_id}")
    records = []

    for config_name in ("ara-ara-corpus", "ara-ara-queries"):
        category = "passages" if config_name.endswith("corpus") else "questions"

        # Both validation and test exist in the MTEB parquet conversion.
        for split_name in ("validation", "test"):
            dataset = load_dataset(dataset_id, config_name, split=split_name)
            for row in dataset:
                records.append(
                    make_record(
                        "MLQA-Arabic",
                        category,
                        row.get("text"),
                        configuration=config_name,
                        group="QA",
                        subject=category,
                    )
                )

    return deduplicate_and_limit(records)


def load_massive_arabic():
    """Short Arabic assistant commands from the MASSIVE NLU benchmark."""
    dataset_id = "mteb/amazon_massive_intent"
    print(f"\nLoading MASSIVE Arabic: {dataset_id}")
    configurations = get_dataset_config_names(dataset_id)
    config_name = "ar" if "ar" in configurations else "default"
    dataset = load_dataset(dataset_id, config_name, split="test")
    records = []

    for row in dataset:
        language = clean_value(row.get("lang"))
        if config_name == "default" and language not in {"ar", "ar-SA"}:
            continue

        intent = clean_value(row.get("label_text") or row.get("label"))
        records.append(
            make_record(
                "MASSIVE-Arabic",
                "assistant-commands",
                row.get("text"),
                configuration=config_name,
                group="spoken-NLU",
                subject=intent or "assistant-command",
                country="Saudi Arabia" if language == "ar-SA" else "",
            )
        )

    return deduplicate_and_limit(records)


def load_arabic_dialects():
    """Regional dialectal/transcribed Arabic covering five dialect groups."""
    dataset_id = "arbml/Arabic_Dialects_Dataset"
    print(f"\nLoading Arabic dialects: {dataset_id}")
    dataset = load_dataset(dataset_id, split="train")
    label_feature = dataset.features.get("label")
    records = []

    for row in dataset:
        label = row.get("label")
        if label_feature is not None and hasattr(label_feature, "int2str"):
            try:
                dialect = label_feature.int2str(label)
            except (TypeError, ValueError):
                dialect = clean_value(label)
        else:
            dialect = clean_value(label)

        records.append(
            make_record(
                "Arabic-Dialects",
                dialect or "dialectal-Arabic",
                row.get("Text") or row.get("text"),
                configuration="default",
                group="dialectal-Arabic",
                subject=dialect or "Unknown dialect",
            )
        )

    return deduplicate_and_limit(records)


def load_all_benchmarks() -> pd.DataFrame:
    records = []
    loaders = (
        load_arabic_mmlu,
        load_mlqa_arabic,
        load_massive_arabic,
        load_arabic_dialects,
    )

    for loader in loaders:
        loaded = loader()
        records.extend(loaded)
        print(f"Loaded {len(loaded):,} unique documents")

    frame = pd.DataFrame(records)
    loaded_names = set(frame["benchmark"].unique())
    missing = EXPECTED_BENCHMARKS - loaded_names
    if missing:
        raise RuntimeError(f"Missing required benchmarks: {sorted(missing)}")

    print("\nBenchmark corpus summary:")
    print(
        frame.groupby("benchmark")
        .agg(documents=("text", "size"), words=("words", "sum"), bytes=("bytes", "sum"))
        .to_string()
    )
    print(f"\nCombined documents: {len(frame):,}")
    print(f"Combined words: {frame['words'].sum():,}")
    print(f"Combined UTF-8 bytes: {frame['bytes'].sum():,}")

    return frame


# ---------------------------------------------------------------------
# Tokenizer wrappers
# ---------------------------------------------------------------------

def encode_in_batches(texts, encode_batch):
    token_lengths = []

    for start in range(0, len(texts), BATCH_SIZE):
        batch = texts[start : start + BATCH_SIZE]
        encoded_batch = encode_batch(batch)

        token_lengths.extend(len(token_ids) for token_ids in encoded_batch)

        completed = min(start + BATCH_SIZE, len(texts))
        print(
            f"\rProcessed {completed:,}/{len(texts):,}",
            end="",
            flush=True,
        )

    print()
    return token_lengths


def benchmark_custom_tokenizer(frame):
    if not MY_TOKENIZER_PATH.exists():
        raise FileNotFoundError(
            f"Your tokenizer was not found: {MY_TOKENIZER_PATH}"
        )

    tokenizer = CoreTokenizer.from_file(str(MY_TOKENIZER_PATH))

    print(f"\nBenchmarking your tokenizer: {MY_TOKENIZER_PATH}")
    print(f"Vocabulary size: {tokenizer.get_vocab_size():,}")

    return encode_in_batches(
        frame["text"].tolist(),
        lambda batch: [
            encoding.ids
            for encoding in tokenizer.encode_batch(
                batch,
                add_special_tokens=False,
            )
        ],
    )


def benchmark_huggingface_tokenizer(name, repository, frame):
    print(f"\nLoading {name}: {repository}")

    load_kwargs = {
        "use_fast": True,
        "trust_remote_code": TRUST_REMOTE_CODE,
        "token": os.environ.get("HF_TOKEN"),
    }
    if name == "Mistral-Small-3.1":
        load_kwargs["fix_mistral_regex"] = True

    try:
        tokenizer = AutoTokenizer.from_pretrained(repository, **load_kwargs)
    except Exception as auto_error:
        # Some new repositories expose tokenizer.json before their tokenizer
        # class is supported by the installed Transformers version. The raw
        # Hugging Face Tokenizers file is sufficient for this compression test.
        try:
            tokenizer_path = hf_hub_download(
                repo_id=repository,
                filename="tokenizer.json",
                token=os.environ.get("HF_TOKEN"),
            )
            core_tokenizer = CoreTokenizer.from_file(tokenizer_path)
            print(f"AutoTokenizer unavailable; using tokenizer.json fallback")
            print(f"Vocabulary size: {core_tokenizer.get_vocab_size():,}")
            print(f"Benchmarking {name}")
            return encode_in_batches(
                frame["text"].tolist(),
                lambda batch: [
                    encoding.ids
                    for encoding in core_tokenizer.encode_batch(
                        batch,
                        add_special_tokens=False,
                    )
                ],
            )
        except Exception as fallback_error:
            raise RuntimeError(
                f"AutoTokenizer failed: {auto_error}; "
                f"tokenizer.json fallback failed: {fallback_error}"
            ) from auto_error

    print(f"Vocabulary size: {len(tokenizer):,}")
    print(f"Benchmarking {name}")

    def encode_batch(batch):
        output = tokenizer(
            batch,
            add_special_tokens=False,
            padding=False,
            truncation=False,
            return_attention_mask=False,
            return_token_type_ids=False,
        )
        return output["input_ids"]

    return encode_in_batches(frame["text"].tolist(), encode_batch)


def benchmark_tiktoken(name, encoding_name, frame):
    print(f"\nLoading {name}: {encoding_name}")

    tokenizer = tiktoken.get_encoding(encoding_name)

    def encode_batch(batch):
        return [
            tokenizer.encode(
                text,
                allowed_special=set(),
                disallowed_special=(),
            )
            for text in batch
        ]

    return encode_in_batches(frame["text"].tolist(), encode_batch)


# ---------------------------------------------------------------------
# Metrics and reports
# ---------------------------------------------------------------------

def build_summary(frame, token_columns, benchmark_name):
    total_documents = len(frame)
    total_words = int(frame["words"].sum())
    total_characters = int(frame["characters"].sum())
    total_bytes = int(frame["bytes"].sum())

    custom_total = int(frame[token_columns["SamAI-Tokenizer"]].sum())
    rows = []

    for tokenizer_name, column in token_columns.items():
        total_tokens = int(frame[column].sum())
        ratio = total_tokens / custom_total

        rows.append(
            {
                "benchmark": benchmark_name,
                "tokenizer": tokenizer_name,
                "documents": total_documents,
                "tokens": total_tokens,
                "tokens_per_word": total_tokens / total_words,
                "characters_per_token": total_characters / total_tokens,
                "bytes_per_token": total_bytes / total_tokens,
                "token_ratio_vs_yours": ratio,
                "extra_tokens_vs_yours_pct": (ratio - 1.0) * 100,
            }
        )

    return pd.DataFrame(rows).sort_values(
        "tokens_per_word",
        ascending=True,
    )


def build_macro_summary(by_benchmark):
    """Equal-weight average so every benchmark contributes exactly 25%."""
    macro = (
        by_benchmark.groupby("tokenizer", as_index=False)
        .agg(
            benchmarks=("benchmark", "nunique"),
            mean_tokens_per_word=("tokens_per_word", "mean"),
            mean_characters_per_token=("characters_per_token", "mean"),
            mean_bytes_per_token=("bytes_per_token", "mean"),
            mean_token_ratio_vs_yours=("token_ratio_vs_yours", "mean"),
            mean_extra_tokens_vs_yours_pct=("extra_tokens_vs_yours_pct", "mean"),
        )
        .sort_values("mean_tokens_per_word")
    )
    macro.insert(0, "rank", range(1, len(macro) + 1))
    return macro


def build_subject_report(frame, token_columns):
    rows = []

    for (benchmark, group, subject), subset in frame.groupby(
        ["benchmark", "group", "subject"],
        dropna=False,
    ):
        words = int(subset["words"].sum())
        byte_count = int(subset["bytes"].sum())
        custom_tokens = int(subset[token_columns["SamAI-Tokenizer"]].sum())

        for tokenizer_name, column in token_columns.items():
            tokens = int(subset[column].sum())

            rows.append(
                {
                    "benchmark": benchmark,
                    "group": group,
                    "subject": subject,
                    "tokenizer": tokenizer_name,
                    "documents": len(subset),
                    "words": words,
                    "tokens": tokens,
                    "tokens_per_word": tokens / words if words else 0,
                    "bytes_per_token": byte_count / tokens if tokens else 0,
                    "token_ratio_vs_yours": (
                        tokens / custom_tokens if custom_tokens else 0
                    ),
                }
            )

    return pd.DataFrame(rows)


# ---------------------------------------------------------------------
# Main benchmark
# ---------------------------------------------------------------------

def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    frame = load_all_benchmarks()
    token_columns = {}

    # Your tokenizer
    column = "tokens__SamAI_Tokenizer"
    frame[column] = benchmark_custom_tokenizer(frame)
    token_columns["SamAI-Tokenizer"] = column

    # Hugging Face competitors
    for name, repository in HF_TOKENIZERS.items():
        try:
            column = f"tokens__{name}"
            frame[column] = benchmark_huggingface_tokenizer(
                name,
                repository,
                frame,
            )
            token_columns[name] = column

        except Exception as error:
            print(f"\nWARNING: Could not benchmark {name}")
            print(f"Repository: {repository}")
            print(f"Reason: {error}")

    # OpenAI tiktoken baselines
    for name, encoding_name in TIKTOKEN_ENCODINGS.items():
        try:
            column = f"tokens__{name}"
            frame[column] = benchmark_tiktoken(
                name,
                encoding_name,
                frame,
            )
            token_columns[name] = column

        except Exception as error:
            print(f"\nWARNING: Could not benchmark {name}: {error}")

    by_benchmark = pd.concat(
        [
            build_summary(subset, token_columns, benchmark_name)
            for benchmark_name, subset in frame.groupby("benchmark", sort=True)
        ],
        ignore_index=True,
    )
    macro_summary = build_macro_summary(by_benchmark)
    overall_summary = build_summary(frame, token_columns, "ALL-POOLED")
    subject_report = build_subject_report(frame, token_columns)

    document_path = OUTPUT_DIR / "arabic_4bench_document_results.csv"
    benchmark_path = OUTPUT_DIR / "arabic_4bench_by_benchmark.csv"
    macro_path = OUTPUT_DIR / "arabic_4bench_macro_average.csv"
    overall_path = OUTPUT_DIR / "arabic_4bench_pooled_overall.csv"
    subject_path = OUTPUT_DIR / "arabic_4bench_by_category.csv"

    frame.to_csv(document_path, index=False)
    by_benchmark.to_csv(benchmark_path, index=False)
    macro_summary.to_csv(macro_path, index=False)
    overall_summary.to_csv(overall_path, index=False)
    subject_report.to_csv(subject_path, index=False)

    print("\n" + "=" * 100)
    print("Four-benchmark Arabic tokenizer comparison: equal-weight macro average")
    print("=" * 100)

    print(
        macro_summary.to_string(
            index=False,
            formatters={
                "mean_tokens_per_word": "{:.4f}".format,
                "mean_characters_per_token": "{:.4f}".format,
                "mean_bytes_per_token": "{:.4f}".format,
                "mean_token_ratio_vs_yours": "{:.4f}".format,
                "mean_extra_tokens_vs_yours_pct": "{:+.2f}%".format,
            },
        )
    )

    print("\n" + "=" * 100)
    print("Results by benchmark")
    print("=" * 100)
    for benchmark_name in sorted(by_benchmark["benchmark"].unique()):
        print(f"\n--- {benchmark_name} ---")
        subset = by_benchmark[by_benchmark["benchmark"] == benchmark_name]
        print(
            subset.to_string(
                index=False,
                formatters={
                    "tokens_per_word": "{:.4f}".format,
                    "characters_per_token": "{:.4f}".format,
                    "bytes_per_token": "{:.4f}".format,
                    "token_ratio_vs_yours": "{:.4f}".format,
                    "extra_tokens_vs_yours_pct": "{:+.2f}%".format,
                },
            )
        )

    print("\nReports saved to:")
    print(macro_path)
    print(benchmark_path)
    print(overall_path)
    print(subject_path)
    print(document_path)


if __name__ == "__main__":
    main()