import os
from pathlib import Path

from datasets import load_dataset
from huggingface_hub import login

from lmr.tokenizer import Tokenizer
from .disk_dataset import DiskDataset
from .proportional_dataset import ProportionalDataset

HF_BASE = "MBZUAI-LLM/SlimPajama-627B-DC"

def initialize_dataset(dataset_config, dataset_dir):
        
    tokenizer = Tokenizer.get_instance()
    
    for component_name, split_token_limits in dataset_config.component_token_limits.items():
        
        if dataset_config.component_whitelist is not None and component_name not in dataset_config.component_whitelist:
            print(f"Component {component_name} not on whitelist, skipping...")
            continue
        
        for split_name in ["train",'validation','test']:
            
            token_limit = split_token_limits[split_name]
            component_dir = dataset_dir / dataset_config.dataset_name / component_name
            
            ds_stream = load_dataset(
                HF_BASE,
                data_files={split_name: f"{split_name}/{component_name}/*.jsonl.zst"},
                split=split_name,
                streaming=True
            )
                        
            output_path = component_dir / f"{split_name}.bin"
            metadata_path = component_dir / f"metadata_{split_name}.json"

            if output_path.exists():
                print(f"[SKIP] {output_path} already exists")
            else:
                print(f"[CREATE] {output_path}")
                DiskDataset.generate_bin(
                    ds_stream,
                    tokenizer,
                    output_path,
                    token_limit=token_limit,
                    metadata_path=metadata_path,
                )