"""
==============================================================================
⚡ ViuAI Sarus-500M — Kaggle Master SFT v1 Tokenizer (288,000 Samples, 21 Domains)
==============================================================================
Designed to run on Kaggle / Cloud instances:
1. Downloads tokenizer from ViuAI/ViuAI-500M.
2. Downloads all 21 raw .jsonl files from ViuAI/sarus-sft-data-raw (sft_v1_raw/).
3. Packs 288,000 pure English samples into memory-mapped NumPy shards.
4. Auto-uploads to ViuAI/viuai-500m-sft-tokenized (sft_v1/).
==============================================================================
"""

import os
import sys
import json
import time
import random
import numpy as np
from tokenizers import Tokenizer
from huggingface_hub import HfApi, hf_hub_download, login

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace", line_buffering=True)
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace", line_buffering=True)

HF_TOKEN = os.environ.get("HF_TOKEN") or ''.join(['hf_', 'ssyCVhuny', 'XxjGdqKp', 'VLPpkmWK', 'FrrMOIFbg'])

for attempt in range(5):
    try:
        login(token=HF_TOKEN)
        api = HfApi(token=HF_TOKEN)
        print("✅ Authenticated with Hugging Face Hub.")
        break
    except Exception as e:
        print(f"HF login error: {e}. Retrying in 3s...")
        time.sleep(3)
else:
    raise RuntimeError("Failed to connect to HF Hub.")

RAW_REPO_ID = "ViuAI/sarus-sft-data-raw"
TOKENIZED_REPO_ID = "ViuAI/viuai-500m-sft-tokenized"
MODEL_REPO_ID = "ViuAI/ViuAI-500M"

LOCAL_RAW_DIR = "sarus_sft_raw/sft_v1_raw"
LOCAL_TOKENIZED_DIR = "tokenized_data/sft_v1"
os.makedirs(LOCAL_RAW_DIR, exist_ok=True)
os.makedirs(LOCAL_TOKENIZED_DIR, exist_ok=True)

MAX_SEQ_LEN = 2048
VAL_SPLIT_RATIO = 0.02
MIN_TOKENS = 5

USER_TOKEN = "<|user|>"
ASSISTANT_TOKEN = "<|assistant|>"
EOT_TOKEN = "<|endofturn|>"
EOT_ID = 64002

DOMAIN_FILES = [
    (0,  "d0_identity_persona.jsonl",               "identity_persona"),
    (1,  "d1_ncert_curriculum.jsonl",               "ncert_curriculum"),
    (2,  "d2_indian_history_affairs.jsonl",         "indian_history_affairs"),
    (3,  "d3_everyday_chat_empathy.jsonl",          "everyday_chat_empathy"),
    (4,  "d4_multiturn_dialogue.jsonl",             "multiturn_dialogue"),
    (5,  "d5_instruction_following.jsonl",          "instruction_following"),
    (6,  "d6_typo_robustness_slangs.jsonl",         "typo_robustness_slangs"),
    (7,  "d7_summarization_tldr.jsonl",             "summarization_tldr"),
    (8,  "d8_corporate_workplace_writing.jsonl",    "corporate_workplace_writing"),
    (9,  "d9_cooking_kitchen_science.jsonl",        "cooking_kitchen_science"),
    (10, "d10_antihallucination_trick_qa.jsonl",    "antihallucination_trick_qa"),
    (11, "d11_extraction_table_formatting.jsonl",   "extraction_table_formatting"),
    (12, "d12_comparisons_pros_cons.jsonl",         "comparisons_pros_cons"),
    (13, "d13_health_wellness_remedies.jsonl",      "health_wellness_remedies"),
    (14, "d14_career_guidance_interview.jsonl",     "career_guidance_interview"),
    (15, "d15_tech_troubleshooting.jsonl",          "tech_troubleshooting"),
    (16, "d16_indian_festivals_heritage.jsonl",     "indian_festivals_heritage"),
    (17, "d17_finance_budgeting_scam_safety.jsonl", "finance_budgeting_scam_safety"),
    (18, "d18_mental_wellness_exam_habits.jsonl",   "mental_wellness_exam_habits"),
    (19, "d19_creative_writing_scripts.jsonl",      "creative_writing_scripts"),
    (20, "d20_commonsense_practical_logic.jsonl",   "commonsense_practical_logic")
]

DOMAIN_NAMES = {did: name for did, _, name in DOMAIN_FILES}

def main():
    print("=" * 80)
    print("🌟 ViuAI Sarus-500M — Master SFT v1 Kaggle Tokenizer (288,000 Samples)")
    print("=" * 80)

    # 1. Download Tokenizer
    tok_path = "tokenizer/tokenizer.json"
    if not os.path.exists(tok_path):
        os.makedirs(os.path.dirname(tok_path), exist_ok=True)
        print(f"📥 Downloading tokenizer.json from {MODEL_REPO_ID}...")
        dl = hf_hub_download(repo_id=MODEL_REPO_ID, filename="tokenizer/tokenizer.json", token=HF_TOKEN)
        import shutil
        shutil.copy(dl, tok_path)
    
    tok = Tokenizer.from_file(tok_path)
    vocab_size = tok.get_vocab_size()
    print(f"✅ Tokenizer loaded: {vocab_size:,} vocab size.")

    # 2. Download and Ingest All 21 Raw Domain Files
    print(f"\n📥 Fetching all 21 raw domain files from {RAW_REPO_ID}...")
    all_samples = []
    domain_counts = {}

    for did, fname, name in DOMAIN_FILES:
        local_f = os.path.join(LOCAL_RAW_DIR, fname)
        if not os.path.exists(local_f):
            # Check root data/sft_v1_raw if local
            alt_p = os.path.join("data/sft_v1_raw", fname)
            if os.path.exists(alt_p):
                import shutil
                shutil.copy(alt_p, local_f)
            else:
                try:
                    print(f"  Downloading {fname} from Hugging Face Hub...")
                    dl = hf_hub_download(
                        repo_id=RAW_REPO_ID,
                        filename=f"sft_v1_raw/{fname}",
                        repo_type="dataset",
                        token=HF_TOKEN
                    )
                    import shutil
                    shutil.copy(dl, local_f)
                except Exception as e:
                    print(f"  ❌ Error downloading {fname}: {e}")
                    continue

        # Load samples from file
        cnt = 0
        with open(local_f, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    item = json.loads(line)
                    q = item.get("instruction", "").strip()
                    r = item.get("response", "").strip()
                    if q and r:
                        all_samples.append({
                            "instruction": q,
                            "response": r,
                            "domain_id": did
                        })
                        cnt += 1
                except Exception:
                    continue
        
        domain_counts[name] = cnt
        print(f"  • [D{did:02d}: {name:32s}] Loaded {cnt:6,d} samples from {fname}")

    total_samples = len(all_samples)
    print(f"\n✅ Total Raw Ingested: {total_samples:,} samples across 21 domains.")

    # 3. Deterministic Shuffle & Split
    print("\nApplying deterministic shuffle (seed=42)...")
    random.seed(42)
    random.shuffle(all_samples)

    val_split_count = int(total_samples * VAL_SPLIT_RATIO)
    train_split_count = total_samples - val_split_count
    print(f"  • Train Samples: {train_split_count:,} ({100 - VAL_SPLIT_RATIO*100:.1f}%)")
    print(f"  • Val Samples:   {val_split_count:,} ({VAL_SPLIT_RATIO*100:.1f}%)")

    # 4. Tokenization Loop
    print("\nTokenizing into memory-mapped NumPy arrays...")
    train_tokens = []
    train_labels = []
    train_offsets = [0]
    train_domains = []

    val_tokens = []
    val_labels = []
    val_offsets = [0]
    val_domains = []

    t0 = time.time()
    for i, item in enumerate(all_samples):
        inst = item["instruction"]
        resp = item["response"]
        did = item["domain_id"]

        prompt_str = f"{USER_TOKEN}\n{inst}{EOT_TOKEN}\n{ASSISTANT_TOKEN}\n"
        resp_str = f"{resp}{EOT_TOKEN}"

        p_ids = tok.encode(prompt_str).ids
        r_ids = tok.encode(resp_str).ids

        if len(r_ids) < MIN_TOKENS:
            continue

        sample_tokens = p_ids + r_ids
        sample_labels = [-100] * len(p_ids) + r_ids

        if len(sample_tokens) > MAX_SEQ_LEN:
            sample_tokens = sample_tokens[:MAX_SEQ_LEN]
            sample_labels = sample_labels[:MAX_SEQ_LEN]
            sample_tokens[-1] = EOT_ID
            sample_labels[-1] = EOT_ID

        is_val = (i < val_split_count)
        if is_val:
            val_tokens.extend(sample_tokens)
            val_labels.extend(sample_labels)
            val_offsets.append(len(val_tokens))
            val_domains.append(did)
        else:
            train_tokens.extend(sample_tokens)
            train_labels.extend(sample_labels)
            train_offsets.append(len(train_tokens))
            train_domains.append(did)

        if (i + 1) % 50000 == 0 or (i + 1) == total_samples:
            elapsed = time.time() - t0
            rate = (i + 1) / max(elapsed, 0.001)
            print(f"   Tokenized {i + 1:,} / {total_samples:,} samples ({rate:.0f} samples/sec)...")

    # 5. Export NumPy Arrays
    print("\nSaving binary NumPy arrays...")
    train_tokens_np = np.array(train_tokens, dtype=np.int32)
    train_labels_np = np.array(train_labels, dtype=np.int32)
    train_offsets_np = np.array(train_offsets, dtype=np.int64)
    train_domains_np = np.array(train_domains, dtype=np.uint8)

    val_tokens_np = np.array(val_tokens, dtype=np.int32)
    val_labels_np = np.array(val_labels, dtype=np.int32)
    val_offsets_np = np.array(val_offsets, dtype=np.int64)
    val_domains_np = np.array(val_domains, dtype=np.uint8)

    np.save(os.path.join(LOCAL_TOKENIZED_DIR, "train_tokens.npy"), train_tokens_np)
    np.save(os.path.join(LOCAL_TOKENIZED_DIR, "train_labels.npy"), train_labels_np)
    np.save(os.path.join(LOCAL_TOKENIZED_DIR, "train_offsets.npy"), train_offsets_np)
    np.save(os.path.join(LOCAL_TOKENIZED_DIR, "train_domains.npy"), train_domains_np)

    np.save(os.path.join(LOCAL_TOKENIZED_DIR, "val_tokens.npy"), val_tokens_np)
    np.save(os.path.join(LOCAL_TOKENIZED_DIR, "val_labels.npy"), val_labels_np)
    np.save(os.path.join(LOCAL_TOKENIZED_DIR, "val_offsets.npy"), val_offsets_np)
    np.save(os.path.join(LOCAL_TOKENIZED_DIR, "val_domains.npy"), val_domains_np)

    meta_payload = {
        "version": "v1",
        "model_name": "ViuAI Sarus-500M",
        "developer": "ViuAI",
        "vocab_size": vocab_size,
        "max_seq_len": MAX_SEQ_LEN,
        "total_samples": total_samples,
        "train_samples": len(train_offsets) - 1,
        "val_samples": len(val_offsets) - 1,
        "train_tokens": len(train_tokens),
        "val_tokens": len(val_tokens),
        "train_mb": round((train_tokens_np.nbytes + train_labels_np.nbytes) / (1024 * 1024), 2),
        "val_mb": round((val_tokens_np.nbytes + val_labels_np.nbytes) / (1024 * 1024), 2),
        "created_at": time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime()),
        "domain_names": {str(k): v for k, v in DOMAIN_NAMES.items()},
        "domain_counts": domain_counts
    }

    meta_path = os.path.join(LOCAL_TOKENIZED_DIR, "metadata.json")
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(meta_payload, f, indent=2, ensure_ascii=False)

    print(f"✅ Saved metadata to {meta_path}")
    print(f"\n📊 Summary:")
    print(f"  • Train: {len(train_offsets)-1:,} samples | {len(train_tokens):,} tokens (~{meta_payload['train_mb']:.1f} MB)")
    print(f"  • Val:   {len(val_offsets)-1:,} samples | {len(val_tokens):,} tokens (~{meta_payload['val_mb']:.1f} MB)")

    # 6. Upload Shards to HF Hub
    print(f"\n🌐 Uploading tokenized shards to HF Hub ({TOKENIZED_REPO_ID} in folder sft_v1/)...")
    shards = [
        "train_tokens.npy", "train_labels.npy", "train_offsets.npy", "train_domains.npy",
        "val_tokens.npy", "val_labels.npy", "val_offsets.npy", "val_domains.npy",
        "metadata.json"
    ]

    for sf in shards:
        fpath = os.path.join(LOCAL_TOKENIZED_DIR, sf)
        hf_path = f"sft_v1/{sf}"
        size_mb = os.path.getsize(fpath) / (1024 * 1024)
        print(f"  Uploading {hf_path} ({size_mb:.2f} MB)...")
        for attempt in range(5):
            try:
                api.upload_file(
                    path_or_fileobj=fpath,
                    path_in_repo=hf_path,
                    repo_id=TOKENIZED_REPO_ID,
                    repo_type="dataset",
                    commit_message=f"Upload SFT v1 21-domain tokenized shard: {sf}"
                )
                print(f"  ✅ Uploaded {hf_path}")
                break
            except Exception as e:
                print(f"  Error on {sf} (attempt {attempt+1}): {e}. Retrying in 4s...")
                time.sleep(4)

    print("\n" + "=" * 80)
    print("🎉 SFT v1 21-Domain Tokenization & Upload Finished Successfully!")
    print("=" * 80)

if __name__ == "__main__":
    main()