#!/usr/bin/env python3
import json, os, sys, torch
from pathlib import Path
from datasets import load_dataset

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from train_ultra import CharTokenizer

DATA_DIR = "/root/cognet-1b/data_1b"
os.makedirs(DATA_DIR, exist_ok=True)

tokenizer = CharTokenizer.load("/root/cognet-1b/tokenizer_v3.json")

def tokenize_texts(texts, desc=""):
    all_ids = []
    for i, text in enumerate(texts):
        if not text or len(text.strip()) < 10:
            continue
        all_ids.extend(tokenizer.encode(text))
        if i % 50000 == 0 and i > 0:
            print(f"  {desc}: {i:,} texts -> {len(all_ids):,} tokens")
    return all_ids

all_token_ids = []

# 1. WIKITEXT-103-RAW
print("1/5 - WikiText-103-RAW...")
try:
    ds = load_dataset("wikitext", "wikitext-103-raw-v1", split="train", trust_remote_code=True)
    texts = [x["text"] for x in ds if x["text"].strip()]
    ids = tokenize_texts(texts, "WikiText-103")
    all_token_ids.extend(ids)
    print(f"  OK WikiText-103: {len(ids):,} tokens")
    del ds, texts
except Exception as e:
    print(f"  FAIL WikiText-103: {e}")

# 2. ALPACA
print("2/5 - Alpaca...")
try:
    ds = load_dataset("tatsu-lab/alpaca", split="train", trust_remote_code=True)
    texts = []
    for x in ds:
        if x["input"]:
            t = f"### Instruction:\n{x['instruction']}\n\n### Input:\n{x['input']}\n\n### Response:\n{x['output']}"
        else:
            t = f"### Instruction:\n{x['instruction']}\n\n### Response:\n{x['output']}"
        texts.append(t)
    ids = tokenize_texts(texts, "Alpaca")
    all_token_ids.extend(ids)
    print(f"  OK Alpaca: {len(ids):,} tokens")
    del ds, texts
except Exception as e:
    print(f"  FAIL Alpaca: {e}")

# 3. OSCAR FRENCH
print("3/5 - OSCAR French...")
try:
    ds = load_dataset("oscar", "unshuffled_deduplicated_fr", split="train", streaming=True, trust_remote_code=True)
    texts = []
    count = 0
    for x in ds:
        texts.append(x["text"])
        count += 1
        if count >= 200000:
            break
    ids = tokenize_texts(texts, "OSCAR-FR")
    all_token_ids.extend(ids)
    print(f"  OK OSCAR-FR: {len(ids):,} tokens")
    del texts
except Exception as e:
    print(f"  FAIL OSCAR-FR: {e}")

# 4. OPEN-ORCA
print("4/5 - OpenOrca...")
try:
    ds = load_dataset("Open-Orca/OpenOrca", split="train", streaming=True, trust_remote_code=True)
    texts = []
    count = 0
    for x in ds:
        t = f"System: {x['system_prompt']}\nHuman: {x['question']}\nAssistant: {x['response']}"
        texts.append(t)
        count += 1
        if count >= 100000:
            break
    ids = tokenize_texts(texts, "OpenOrca")
    all_token_ids.extend(ids)
    print(f"  OK OpenOrca: {len(ids):,} tokens")
    del texts
except Exception as e:
    print(f"  FAIL OpenOrca: {e}")

# 5. CODE PYTHON
print("5/5 - Code Python...")
try:
    ds = load_dataset("bigcode/the-stack", data_dir="data/python", split="train", streaming=True, trust_remote_code=True)
    texts = []
    count = 0
    for x in ds:
        texts.append(x["content"])
        count += 1
        if count >= 50000:
            break
    ids = tokenize_texts(texts, "Code-Python")
    all_token_ids.extend(ids)
    print(f"  OK Code-Python: {len(ids):,} tokens")
    del texts
except Exception as e:
    print(f"  FAIL Code: {e}")

# LOCAL AICL
print("Adding local AICL data...")
import glob as glob_mod
local_ids = []
data_dir = "/root/cognet-1b/data"
for filepath in sorted(glob_mod.glob(os.path.join(data_dir, "*.json")) + glob_mod.glob(os.path.join(data_dir, "*.jsonl"))):
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            content = f.read()
        try:
            data = json.loads(content)
            if isinstance(data, dict):
                for v in data.values():
                    if isinstance(v, str) and len(v) > 20:
                        local_ids.extend(tokenizer.encode(v))
                    elif isinstance(v, list):
                        for item in v:
                            if isinstance(item, str) and len(item) > 20:
                                local_ids.extend(tokenizer.encode(item))
                            elif isinstance(item, dict):
                                for iv in item.values():
                                    if isinstance(iv, str) and len(iv) > 20:
                                        local_ids.extend(tokenizer.encode(iv))
        except json.JSONDecodeError:
            if len(content) > 20:
                local_ids.extend(tokenizer.encode(content))
    except:
        pass
all_token_ids.extend(local_ids)
print(f"  OK Local AICL: {len(local_ids):,} tokens")

# SAVE
print(f"TOTAL TOKENS: {len(all_token_ids):,}")
tokens = torch.tensor(all_token_ids, dtype=torch.long)
split = int(len(tokens) * 0.95)
train_tokens = tokens[:split]
val_tokens = tokens[split:]
torch.save(train_tokens, os.path.join(DATA_DIR, "train_tokens.pt"))
torch.save(val_tokens, os.path.join(DATA_DIR, "val_tokens.pt"))
print(f"Train: {len(train_tokens):,} tokens ({len(train_tokens)/1e6:.1f}M)")
print(f"Val:   {len(val_tokens):,} tokens ({len(val_tokens)/1e6:.1f}M)")
print("DATA DOWNLOAD COMPLETE!")