"""
prepare_data.py — скачивает датасет и готовит .txt файлы для обучения.
Использует датасет текстов песен с HuggingFace (не требует токена).
"""

import os
from pathlib import Path

TRAIN_DIR = Path("data/train")
VAL_DIR = Path("data/val")

TRAIN_DIR.mkdir(parents=True, exist_ok=True)
VAL_DIR.mkdir(parents=True, exist_ok=True)

# Если данные уже есть — пропускаем
train_files = list(TRAIN_DIR.glob("*.txt"))
if train_files:
    print(f"[Data] Данные уже есть ({len(train_files)} файлов). Пропускаем скачивание.")
    exit(0)

print("[Data] Скачиваем датасет...")

try:
    from datasets import load_dataset

    # Датасет текстов песен — английский, публичный, без авторизации
    # ~400k песен, каждая со словами
    ds = load_dataset(
        "MB24/genius-lyrics",  # публичный датасет текстов песен
        split="train",
        trust_remote_code=True,
    )

    print(f"[Data] Загружено {len(ds)} записей")

    # Берём первые 50k для обучения, 5k для валидации
    train_texts = []
    val_texts = []

    for i, item in enumerate(ds):
        # Поле с текстом песни
        lyrics = item.get("lyrics") or item.get("text") or item.get("content") or ""
        if not lyrics or len(lyrics) < 50:
            continue

        # Чистим текст
        text = lyrics.strip() + "\n\n"

        if i % 10 == 0:
            val_texts.append(text)
        else:
            train_texts.append(text)

        if len(train_texts) >= 50000:
            break

    print(f"[Data] Train: {len(train_texts)}, Val: {len(val_texts)}")

    # Сохраняем чанками по 5000 записей в файл
    chunk_size = 5000
    for chunk_i, start in enumerate(range(0, len(train_texts), chunk_size)):
        chunk = train_texts[start:start + chunk_size]
        out_path = TRAIN_DIR / f"genius_train_{chunk_i:03d}.txt"
        out_path.write_text("".join(chunk), encoding="utf-8")
        print(f"[Data] Записан {out_path} ({len(chunk)} текстов)")

    val_path = VAL_DIR / "genius_val.txt"
    val_path.write_text("".join(val_texts[:1000]), encoding="utf-8")
    print(f"[Data] Записан {val_path}")

except Exception as e:
    print(f"[Data] Ошибка загрузки датасета: {e}")
    print("[Data] Пробуем альтернативный датасет...")

    try:
        from datasets import load_dataset

        # Запасной вариант — WikiText на русском/английском
        ds = load_dataset("wikimedia/wikipedia", "20231101.ru", split="train", streaming=True, trust_remote_code=True)

        train_texts = []
        val_texts = []

        for i, item in enumerate(ds):
            text = item.get("text", "").strip()
            if len(text) < 100:
                continue
            text = text[:2000] + "\n\n"  # берём первые 2000 символов статьи

            if i % 10 == 0:
                val_texts.append(text)
            else:
                train_texts.append(text)

            if len(train_texts) >= 20000:
                break

        chunk_size = 5000
        for chunk_i, start in enumerate(range(0, len(train_texts), chunk_size)):
            chunk = train_texts[start:start + chunk_size]
            out_path = TRAIN_DIR / f"wiki_train_{chunk_i:03d}.txt"
            out_path.write_text("".join(chunk), encoding="utf-8")

        val_path = VAL_DIR / "wiki_val.txt"
        val_path.write_text("".join(val_texts[:500]), encoding="utf-8")
        print(f"[Data] Wikipedia датасет загружен успешно")

    except Exception as e2:
        print(f"[Data] Ошибка запасного датасета: {e2}")
        print("[Data] Создаём минимальные демо-данные...")

        demo = (
            "Языковая модель обучается на текстах.\n"
            "Нейронные сети — основа современного ИИ.\n"
            "Трансформеры используют механизм внимания.\n"
            "Python — язык для машинного обучения.\n"
        ) * 500

        (TRAIN_DIR / "demo.txt").write_text(demo, encoding="utf-8")
        (VAL_DIR / "demo_val.txt").write_text(demo[:5000], encoding="utf-8")
        print("[Data] Демо-данные созданы.")

print("[Data] Подготовка данных завершена.")