import os
import numpy as np
import tiktoken
from datasets import load_dataset
from tqdm import tqdm

# --- KONFIGURATION ---
OUTPUT_DIR = "data/alpaca_cleaned_mixed_NEW"
# Wichtig: Das muss exakt der Tokenizer sein, den dein Modell verwendet (meist GPT-2)
TOKENIZER_NAME = "gpt2" 
SEED = 1337

# Balance: Wie viel Alpaca vs. FineWeb?
# Zu viel FineWeb = Modell antwortet nicht im Chat-Stil
# Zu wenig FineWeb = Modell wird dumm (vergisst Weltwissen)
FINEWEB_SAMPLES = 520000

enc = tiktoken.get_encoding(TOKENIZER_NAME)
EOS_TOKEN = "<|endoftext|>" # End of Sequence Token

def format_prompt_with_mask(instruction, input_text, output):
    """
    Formatiert den Prompt und erstellt die Loss-Maske.
    Format:
    Instruction: ...
    Input: ... (optional)
    Response: ... <|endoftext|>
    """
    # 1. Den "Prompt"-Teil bauen (Frage) -> Wird NICHT trainiert (Maske 0)
    if input_text and input_text.strip():
        prompt_text = f"Instruction:\n{instruction}\n\nInput:\n{input_text}\n\nResponse:\n"
    else:
        prompt_text = f"Instruction:\n{instruction}\n\nResponse:\n"
    
    # 2. Den "Completion"-Teil bauen (Antwort) -> Wird TRAINIERT (Maske 1)
    completion_text = f"{output}{EOS_TOKEN}"
    
    # 3. Tokenisieren
    # encode_plain verhindert, dass special tokens im normalen Text interpretiert werden
    prompt_ids = enc.encode(prompt_text, allowed_special={'<|endoftext|>'})
    completion_ids = enc.encode(completion_text, allowed_special={'<|endoftext|>'})
    
    # 4. Zusammenfügen
    full_ids = prompt_ids + completion_ids
    
    # 5. Maske erstellen
    # 0 = Ignorieren (Loss wird hier nicht berechnet)
    # 1 = Trainieren (Modell soll lernen, das vorherzusagen)
    mask = [0] * len(prompt_ids) + [1] * len(completion_ids)
    
    return full_ids, mask

def main():
    np.random.seed(SEED)
    print(f"🚀 Starte Prepare-Script für SmaLLMPro (350M SFT)...")
    print(f"📚 Tokenizer: {TOKENIZER_NAME}")
    
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # --- 1. DATENSÄTZE LADEN ---
    print("📥 Lade 'yahma/alpaca-cleaned' (Chat-Instruktionen)...")
    alpaca = load_dataset("yahma/alpaca-cleaned", split='train')
    
    print(f"📥 Lade 'HuggingFaceFW/fineweb-edu' (Sample-10BT) für {FINEWEB_SAMPLES} Samples...")
    fineweb = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split='train', streaming=True)
    
    all_tokens = []
    all_masks = []
    
    # --- 2. ALPACA VERARBEITEN (Masking aktiv) ---
    print("⚙️  Verarbeite Alpaca...")
    for ex in tqdm(alpaca, desc="Alpaca"):
        ids, mask = format_prompt_with_mask(ex['instruction'], ex['input'], ex['output'])
        all_tokens.extend(ids)
        all_masks.extend(mask)
        
    alpaca_len = len(all_tokens)
    print(f"   -> Alpaca Tokens: {alpaca_len:,}")

    # --- 3. FINEWEB VERARBEITEN (Wissenserhalt) ---
    # Hier setzen wir die Maske auf 1 für den GANZEN Text.
    # Warum? Das Modell soll das Weltwissen aktiv auffrischen, nicht nur als Prompt sehen.
    print("⚙️  Verarbeite FineWeb (Anti-Forgetting)...")
    fw_iter = iter(fineweb)
    fw_count = 0
    fw_tokens_count = 0
    
    for _ in tqdm(range(FINEWEB_SAMPLES), desc="FineWeb"):
        try:
            ex = next(fw_iter)
            text = ex['text'] + EOS_TOKEN
            ids = enc.encode(text, allowed_special={EOS_TOKEN})
            
            all_tokens.extend(ids)
            # Alles lernen!
            all_masks.extend([1] * len(ids)) 
            
            fw_tokens_count += len(ids)
            fw_count += 1
        except StopIteration:
            break
            
    print(f"   -> FineWeb Tokens: {fw_tokens_count:,} (aus {fw_count} Dokumenten)")

    # --- 4. SPEICHERN ---
    total_tokens = len(all_tokens)
    print(f"\n💾 Speichere {total_tokens:,} Tokens in '{OUTPUT_DIR}'...")
    
    # Tokens als uint16 (spart Platz, reicht für GPT-2 Vocab ~50k)
    token_arr = np.array(all_tokens, dtype=np.uint16)
    token_arr.tofile(os.path.join(OUTPUT_DIR, "train.bin"))
    
    # Maske als uint8 (braucht nur 0 oder 1)
    mask_arr = np.array(all_masks, dtype=np.uint8)
    mask_arr.tofile(os.path.join(OUTPUT_DIR, "train_mask.bin"))
    
    # --- 5. SANITY CHECK (Ganz wichtig!) ---
    print("\n🔍 --- SANITY CHECK ---")
    print("Ich dekodiere die ersten 50 Tokens des ersten Beispiels, um zu prüfen, ob alles stimmt.")
    print("Grün (TRAIN) = Was das Modell lernt. Grau (IGNORE) = Was das Modell nur liest.")
    
    check_len = 100
    sample_ids = all_tokens[:check_len]
    sample_mask = all_masks[:check_len]
    
    # Wir rekonstruieren den Text und zeigen, was maskiert ist
    decoded_parts = []
    for t_id, m_val in zip(sample_ids, sample_mask):
        token_str = enc.decode([t_id])
        if m_val == 1:
            decoded_parts.append(f"\033[92m{token_str}\033[0m") # Grün für Training
        else:
            decoded_parts.append(f"\033[90m{token_str}\033[0m") # Grau für Prompt
            
    print("".join(decoded_parts))
    print("\n(Legende: \033[90mGrau=Prompt/Ignoriert\033[0m, \033[92mGrün=Response/Gelernt\033[0m)")
    
    if len(token_arr) != len(mask_arr):
        print("\n❌ ACHTUNG: Token und Mask Array sind unterschiedlich lang! Irgendwas stimmt nicht!")
    else:
        print("\n✅ Alles perfekt. Arrays sind synchron. Du kannst trainieren.")

if __name__ == "__main__":
    main()