import os
import random
from pathlib import Path
from datasets import load_dataset

GUTENBERG = "D:/gutenberg_en.txt"
RU_BOOKS  = "D:/ru_books.txt"

SCRIPT_DIR = Path(__file__).parent
OUT_DIR = SCRIPT_DIR / "sample"
OUT_DIR.mkdir(parents=True, exist_ok=True)

OUT_EN = OUT_DIR / "gutenberg_sample.txt"
OUT_RU = OUT_DIR / "ru_books_sample.txt"

N_SAMPLES = 5000
MIN_LEN   = 100
MAX_LEN   = 2000
MAX_SCAN  = 2_000_000


def download_gutenberg():
    print("downloading gutenberg...")
    ds = load_dataset("AdhyanshVerma/pg-en", split="train", streaming=True)
    with open(GUTENBERG, "w", encoding="utf-8") as f:
        n = 0
        for s in ds:
            t = s.get("text", "").strip()
            if t:
                f.write(t + "\n")
                n += 1
            if n >= 500_000:
                break
    print(f"gutenberg: {n} lines")


def download_ru_books():
    print("downloading ru books...")
    ds = load_dataset("maxzt/RuHeritage-Corpus", split="train", streaming=True)
    with open(RU_BOOKS, "w", encoding="utf-8") as f:
        n = 0
        for s in ds:
            t = s.get("text", "").strip()
            if t:
                f.write(t + "\n")
                n += 1
            if n >= 500_000:
                break
    print(f"ru books: {n} lines")


def make_sample(src_path, dst_path, n_samples, min_len, max_len):
    if not os.path.exists(src_path):
        print(f"missing: {src_path}")
        return 0

    candidates = []
    total = 0

    with open(src_path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            total += 1
            line = line.strip()
            if min_len <= len(line) <= max_len:
                candidates.append(line)
            if total >= MAX_SCAN:
                break

    if not candidates:
        return 0

    random.shuffle(candidates)
    selected = candidates[:n_samples]

    with open(dst_path, "w", encoding="utf-8") as f:
        for line in selected:
            f.write(line + "\n")

    return len(selected)


if __name__ == "__main__":
    random.seed(42)

    if not os.path.exists(GUTENBERG):
        download_gutenberg()

    if not os.path.exists(RU_BOOKS):
        download_ru_books()

    en = make_sample(GUTENBERG, OUT_EN, N_SAMPLES, MIN_LEN, MAX_LEN)
    ru = make_sample(RU_BOOKS, OUT_RU, N_SAMPLES, MIN_LEN, MAX_LEN)

    print(f"en: {en}")
    print(f"ru: {ru}")