import json
import argparse
import multiprocessing as mp
from faker import Faker
import os
import random

def generate_chunk(chunk_id, num_examples, languages, output_dir, chunk_size=10000):
    """
    Генерирует чанк датасета.
    Если languages — список, то для каждого примера случайно выбирается язык.
    """
    fakers = {lang: Faker(lang) for lang in languages}
    filename = os.path.join(output_dir, f"chunk_{chunk_id:05d}.jsonl")
    with open(filename, 'w', encoding='utf-8') as f:
        for _ in range(num_examples):
            lang = random.choice(languages)
            fake = fakers[lang]
            text = fake.paragraph(nb_sentences=random.randint(5, 15), variable_nb_sentences=True)
            record = {"text": text, "language": lang}
            f.write(json.dumps(record, ensure_ascii=False) + '\n')
    return filename

def main():
    parser = argparse.ArgumentParser(description="Генерация мультиязычного синтетического датасета в формате JSONL")
    parser.add_argument("--total_examples", type=int, default=1_000_000,
                        help="Общее количество примеров (текстов) для генерации")
    parser.add_argument("--languages", type=str, nargs='+', default=["en_US", "ru_RU", "zh_CN"],
                        help="Список локалей через пробел (например, en_US ru_RU zh_CN)")
    parser.add_argument("--output_dir", type=str, default="./multilingual_dataset",
                        help="Директория для сохранения чанков")
    parser.add_argument("--chunk_size", type=int, default=10000,
                        help="Количество примеров в одном файле-чанке")
    parser.add_argument("--num_workers", type=int, default=4,
                        help="Количество параллельных процессов для генерации")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    num_chunks = (args.total_examples + args.chunk_size - 1) // args.chunk_size
    chunks = []
    for i in range(num_chunks):
        start = i * args.chunk_size
        end = min(start + args.chunk_size, args.total_examples)
        num_in_chunk = end - start
        chunks.append((i, num_in_chunk, args.languages, args.output_dir, args.chunk_size))

    pool = mp.Pool(processes=args.num_workers)
    results = []
    for chunk in chunks:
        results.append(pool.apply_async(generate_chunk, chunk))
    pool.close()
    pool.join()

    for r in results:
        r.get()

    print(f"Генерация завершена. Чанки сохранены в {args.output_dir}")

if __name__ == "main":
    main()