"""
Script to download and prepare WikiText dataset for training.
WikiText is a common benchmark dataset for language modeling.
"""

import os
import urllib.request
import zipfile
import shutil


def download_wikitext(output_dir='./data'):
    """Download WikiText-2 dataset."""
    os.makedirs(output_dir, exist_ok=True)
    
    url = "https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip"
    zip_path = os.path.join(output_dir, "wikitext-2-v1.zip")
    
    print("Downloading WikiText-2 dataset...")
    if not os.path.exists(zip_path):
        urllib.request.urlretrieve(url, zip_path)
        print(f"Downloaded to {zip_path}")
    else:
        print(f"File already exists: {zip_path}")
    
    print("Extracting...")
    extract_dir = os.path.join(output_dir, "wikitext-2")
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(output_dir)
    
    print("Processing files...")
    
    # Read and combine train/valid/test files
    splits = ['train', 'valid', 'test']
    for split in splits:
        input_file = os.path.join(extract_dir, f"wiki.{split}.tokens")
        output_file = os.path.join(output_dir, f"wikitext_{split}.txt")
        
        if os.path.exists(input_file):
            with open(input_file, 'r', encoding='utf-8') as f_in:
                with open(output_file, 'w', encoding='utf-8') as f_out:
                    f_out.write(f_in.read())
            print(f"Created {output_file}")
    
    print(f"\nWikiText-2 dataset ready in {output_dir}/")
    print("Files:")
    print(f"  - wikitext_train.txt (training data)")
    print(f"  - wikitext_valid.txt (validation data)")
    print(f"  - wikitext_test.txt (test data)")


if __name__ == '__main__':
    download_wikitext()