"""
Production Cloud Tokenizer Training Runner - ViuMini-MoE-242M
=============================================================
Architecture: Byte-Level BPE (48,000 Vocabulary)
Pre-Tokenizer: Llama-3 / Indic Regex with Unicode Combining Marks (\p{M})
Data Corpus: ViuAI/viu-mini-raw-pretrain (>105B Tokens Pretraining Corpus)
Ratio: 45% Devanagari Hindi + 35% Hinglish + 20% English & Reasoning
Special Tokens: 13 tokens including <soch>, </soch>, <|user|>, <|assistant|>

Designed for high-speed execution in Google Colab (CPU or T4) / Kaggle Notebooks.
Execution time: ~4-5 minutes on cloud 10 Gbps network.
"""

import os
import sys
import time
from pathlib import Path

# Ensure dependencies are available
try:
    import datasets
    import tokenizers
    import huggingface_hub
except ImportError:
    print("[setup] Installing required cloud dependencies...")
    os.system("pip install -q tokenizers datasets huggingface_hub pyarrow")

from huggingface_hub import HfApi
from tokenizers import Tokenizer, Regex
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Split, ByteLevel, Sequence
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.normalizers import NFC, Sequence as NormSequence
from datasets import load_dataset


# Configuration
HF_DATASET_REPO = "ViuAI/viu-mini-raw-pretrain"
VOCAB_SIZE = 48000
TOTAL_TRAIN_LINES = 500000
OUTPUT_TOKENIZER_PATH = "tokenizer.json"

def _get_hf_token():
    # SECURITY: token kabhi hardcode mat karo — env / Kaggle / Colab secrets se lo.
    tok = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if not tok:
        try:
            from kaggle_secrets import UserSecretsClient
            tok = UserSecretsClient().get_secret("HF_TOKEN")
        except Exception:
            pass
    if not tok:
        try:
            from google.colab import userdata
            tok = userdata.get("HF_TOKEN")
        except Exception:
            pass
    return tok
HF_TOKEN = _get_hf_token()

# Llama 3 / Indic regex pattern: includes \p{M} so Devanagari vowel signs stay attached to consonants
INDIC_LLAMA3_PATTERN = (
    r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|"
    r"[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+|"
    r"\p{N}{1,3}|"
    r" ?[^\s\p{L}\p{N}\p{M}]+[\r\n]*|"
    r"\s*[\r\n]+|"
    r"\s+(?!\S)|\s+"
)

SPECIAL_TOKENS = [
    "<pad>",
    "<bos>",
    "<eos>",
    "<unk>",
    "<mask>",  # v2 (2026-09-24): `<mask|>` typo fixed
    "<|hindi|>",
    "<|english|>",
    "<|hinglish|>",
    "<soch>",
    "</soch>",
    "<|user|>",
    "<|assistant|>",
    "<|system|>",
]

BENCHMARK_SAMPLES = [
    # Hinglish (Romanized Hindi)
    "mai tumse bahut pyaar karta hu",
    "bhai kal party me kya scene hai?",
    "yaar ye phone ka network bahut slow hai",
    "tumne khana khaya kya abhi tak?",
    "<|user|> mujhe train ka status batao\n<|assistant|> <soch> Checking train schedule... </soch> aapki train time par hai.",
    # Devanagari Hindi
    "नमस्ते, आप कैसे हैं?",
    "मुझे हिंदी में कहानी सुनाओ",
    "भारतीय संविधान का अनुच्छेद 21 जीवन के अधिकार की रक्षा करता है।",
    "क्षत्रिय ज्ञानी व्यक्ति त्रिशूल लेकर आया।",
    "<|user|> भारत की राजधानी क्या है?\n<|assistant|> <soch> नई दिल्ली भारत की राजधानी है। </soch> नई दिल्ली।",
    # English
    "The quick brown fox jumps over the lazy dog.",
    "Artificial intelligence is transforming real-world problem solving.",
    "Can you explain quantum computing in simple terms?",
    "DeepSeek-R1 utilizes reinforcement learning for reasoning verification.",
]


def build_cloud_streaming_iterator(repo_id: str, max_lines: int, token: str):
    """Stream lines from Hub. v2 mix (OPTIONAL rule): 30% Hindi, 35% Hinglish, 35% English (English-boosted; v1 was 45/35/20)."""
    api = HfApi(token=token)
    repo_files = api.list_repo_files(repo_id, repo_type="dataset")

    hindi_files = [f for f in repo_files if (f.startswith("hindi/") or f.startswith("hindi_fixed/")) and f.endswith(".parquet")]
    hinglish_files = [f for f in repo_files if f.startswith("hinglish/") and (f.endswith(".parquet") or f.endswith(".jsonl"))]
    english_files = [f for f in repo_files if (f.startswith("distilled/") or f.startswith("english_fixed/")) and f.endswith(".parquet")]

    print(f"[corpus] Discovered on Hub:")
    print(f"  - Hindi files:    {len(hindi_files)}")
    print(f"  - Hinglish files: {len(hinglish_files)}")
    print(f"  - English files:  {len(english_files)}")

    target_hinglish = int(max_lines * 0.35)
    target_hindi = int(max_lines * 0.30)
    target_english = max_lines - target_hindi - target_hinglish

    print(f"[targets] Hindi: {target_hindi:,} lines | Hinglish: {target_hinglish:,} lines | English: {target_english:,} lines")

    def fetch_stream(file_list, target_count, category_name):
        emitted = 0
        for fpath in file_list:
            if emitted >= target_count:
                break
            try:
                ds = load_dataset(repo_id, data_files=fpath, split="train", streaming=True, token=token)
                for row in ds:
                    text = (row.get("text", "") or "").strip()
                    if len(text) >= 15:
                        yield text
                        emitted += 1
                        if emitted >= target_count:
                            break
            except Exception as exc:
                continue
        print(f"[stream] Finished streaming {emitted:,} lines for {category_name}")

    def interleave_generator():
        hi_gen = fetch_stream(hindi_files, target_hindi, "Hindi")
        h_gen = fetch_stream(hinglish_files, target_hinglish, "Hinglish")
        e_gen = fetch_stream(english_files, target_english, "English/Reasoning")

        generators = {"hindi": hi_gen, "hinglish": h_gen, "english": e_gen}
        pattern = ["hinglish"] * 7 + ["hindi"] * 6 + ["english"] * 7  # v2: 35/30/35 (v1 tha 5/3/2 = 50/30/20)
        active = {"hindi": True, "hinglish": True, "english": True}

        total_emitted = 0
        while total_emitted < max_lines and any(active.values()):
            progress = False
            for cat in pattern:
                if not active[cat]:
                    continue
                try:
                    line = next(generators[cat])
                    yield line
                    total_emitted += 1
                    progress = True
                    if total_emitted % 50000 == 0:
                        print(f"  [progress] Streamed {total_emitted:,} / {max_lines:,} training lines...")
                    if total_emitted >= max_lines:
                        return
                except StopIteration:
                    active[cat] = False
            if not progress:
                break

    return interleave_generator()


def run_evaluation(tok: Tokenizer):
    """Run comprehensive quality audit on the trained tokenizer."""
    print("\n" + "=" * 70)
    print("PRODUCTION TOKENIZER EVALUATION & QUALITY AUDIT")
    print("=" * 70)

    # 1. Special token isolation check
    print("\n[Audit 1] Special Tokens Atomic Isolation Check:")
    isolation_passed = True
    for st in SPECIAL_TOKENS:
        enc = tok.encode(st)
        if len(enc.tokens) == 1 and enc.tokens[0] == st:
            print(f"  [PASS] {st:15} -> Token ID: {enc.ids[0]}")
        else:
            isolation_passed = False
            print(f"  [FAIL] {st:15} -> Split into {enc.tokens} (IDs: {enc.ids})")

    if isolation_passed:
        print("  -> All 13 special tokens successfully isolated as single atomic IDs.")

    # 2. Multi-lingual fertility benchmark
    print("\n[Audit 2] Multi-Lingual Fertility & Compression Benchmark:")
    category_fertility = {"Hinglish": [], "Hindi": [], "English": []}

    for sample in BENCHMARK_SAMPLES:
        enc = tok.encode(sample)
        words = sample.split()
        ratio = len(enc.tokens) / max(len(words), 1)

        if any(ord(c) >= 0x0900 and ord(c) <= 0x097F for c in sample):
            category_fertility["Hindi"].append(ratio)
        elif "mai" in sample or "bhai" in sample or "yaar" in sample or "train" in sample:
            category_fertility["Hinglish"].append(ratio)
        else:
            category_fertility["English"].append(ratio)

    print("\nFertility Evaluation Summary (Tokens per Word):")
    for category, ratios in category_fertility.items():
        if ratios:
            avg_ratio = sum(ratios) / len(ratios)
            target = 1.8 if category == "Hindi" else (1.5 if category == "Hinglish" else 1.3)
            status = "EXCELLENT" if avg_ratio <= target else ("ACCEPTABLE" if avg_ratio <= target + 0.3 else "HIGH")
            print(f"  - {category:10}: {avg_ratio:.2f} tokens/word (Target: <= {target:.1f}) [{status}]")

    # 3. Round-trip lossless decoding check
    print("\n[Audit 3] Round-Trip Lossless Decoding Check:")
    fidelity_passed = True
    for sample in BENCHMARK_SAMPLES:
        enc = tok.encode(sample)
        decoded = tok.decode(enc.ids, skip_special_tokens=False)
        if decoded.strip() != sample.strip():
            fidelity_passed = False
            print(f"  [FAIL] Round-trip mismatch:\n    Expected: {sample}\n    Actual:   {decoded}")

    if fidelity_passed:
        print("  [PASS] 100% Round-trip lossless fidelity verified across all test samples.")

    print("=" * 70 + "\n")
    return isolation_passed and fidelity_passed


def main():
    start_time = time.time()
    print("=" * 70)
    print("ViuMini-MoE-242M: Production Tokenizer Training (Indic-Optimized)")
    print(f"Target Vocabulary: {VOCAB_SIZE:,} | Training Samples: {TOTAL_TRAIN_LINES:,}")
    print(f"Corpus Source: {HF_DATASET_REPO}")
    print("=" * 70)

    # Initialize Tokenizer Architecture with Indic-optimized Llama-3 Split Regex
    tok = Tokenizer(BPE(unk_token="<unk>"))
    tok.normalizer = NormSequence([NFC()])
    tok.pre_tokenizer = Sequence([
        Split(pattern=Regex(INDIC_LLAMA3_PATTERN), behavior="isolated"),
        ByteLevel(add_prefix_space=False, use_regex=False),
    ])
    tok.decoder = ByteLevelDecoder()

    trainer = BpeTrainer(
        vocab_size=VOCAB_SIZE,
        min_frequency=2,
        show_progress=True,
        special_tokens=SPECIAL_TOKENS,
        initial_alphabet=ByteLevel.alphabet(),
    )

    # Build streaming training iterator
    print("[train] Preparing balanced multi-lingual streaming iterator...")
    training_iterator = build_cloud_streaming_iterator(
        repo_id=HF_DATASET_REPO,
        max_lines=TOTAL_TRAIN_LINES,
        token=HF_TOKEN,
    )

    # Train Byte-Level BPE
    print(f"[train] Commencing Rust-accelerated BPE training on {TOTAL_TRAIN_LINES:,} samples...")
    tok.train_from_iterator(training_iterator, trainer=trainer, length=TOTAL_TRAIN_LINES)

    # Save trained tokenizer locally
    tok.save(OUTPUT_TOKENIZER_PATH)
    file_size_mb = os.path.getsize(OUTPUT_TOKENIZER_PATH) / (1024 * 1024)
    print(f"[train] Tokenizer saved locally: {OUTPUT_TOKENIZER_PATH} ({file_size_mb:.2f} MB)")

    # Execute Evaluation Suite
    audit_success = run_evaluation(tok)

    # Upload to Hugging Face Hub
    if HF_TOKEN:
        print(f"[hub] Uploading production tokenizer to {HF_DATASET_REPO}...")
        api = HfApi(token=HF_TOKEN)
        api.upload_file(
            path_or_fileobj=OUTPUT_TOKENIZER_PATH,
            path_in_repo="tokenizer/tokenizer.json",
            repo_id=HF_DATASET_REPO,
            repo_type="dataset",
        )
        print(f"[hub] Successfully uploaded tokenizer/tokenizer.json to {HF_DATASET_REPO}!")

    elapsed = time.time() - start_time
    print(f"[complete] All steps finished successfully in {elapsed / 60:.2f} minutes.")


if __name__ == "__main__":
    main()