"""Factory for creating datasets based on configuration."""

from typing import Union, Optional
from taoTrain.config import TrainingConfig, TrainingModeEnum
from taoTrain.data.hf_pretrain import PretrainDataset
from taoTrain.data.hf_sft import SFTDataset
from taoTrain.data.hf_rl import RLDataset
from taoTrain.data.pretrain_jsonl import PretrainJSONLDataset
from taoTrain.data.sft_jsonl import SFTJSONLDataset
from taoTrain.data.rl_jsonl import RLJSONLDataset
from taoTrain.data.vlm_jsonl import VLMJSONLDataset, VLMSFTJSONLDataset


class DatasetFactory:
    """Factory for creating datasets based on configuration."""
    
    # Registry of dataset classes by mode and backend
    DATASETS = {
        (TrainingModeEnum.PRETRAIN, "huggingface"): PretrainDataset,
        (TrainingModeEnum.SFT, "huggingface"): SFTDataset,
        (TrainingModeEnum.RL, "huggingface"): RLDataset,
        (TrainingModeEnum.PRETRAIN, "jsonl"): PretrainJSONLDataset,
        (TrainingModeEnum.SFT, "jsonl"): SFTJSONLDataset,
        (TrainingModeEnum.RL, "jsonl"): RLJSONLDataset,
        (TrainingModeEnum.VLM, "jsonl"): VLMJSONLDataset,
        (TrainingModeEnum.VLM_SFT, "jsonl"): VLMSFTJSONLDataset,
    }
    
    @staticmethod
    def create_dataset(
        config: TrainingConfig,
        split: str = "train",
    ) -> Union[PretrainDataset, SFTDataset, RLDataset, PretrainJSONLDataset, SFTJSONLDataset, RLJSONLDataset]:
        """
        Create dataset instance based on configuration.
        
        Args:
            config: Training configuration
            split: Dataset split (train, validation, test) - primarily for HuggingFace datasets
        
        Returns:
            Dataset instance matching the configured mode and backend
        
        Raises:
            ValueError: If configuration is invalid or unsupported mode/backend combination
        """
        # Determine backend: JSONL or HuggingFace
        if config.dataset.local:
            backend = "jsonl"
        else:
            backend = "huggingface"
        
        # Get mode
        mode = config.mode
        
        # Look up dataset class
        key = (mode, backend)
        if key not in DatasetFactory.DATASETS:
            raise ValueError(
                f"Unsupported dataset configuration: mode={mode.value}, backend={backend}. "
                f"Supported: {list(DatasetFactory.DATASETS.keys())}"
            )
        
        dataset_class = DatasetFactory.DATASETS[key]
        
        # Instantiate dataset
        if backend == "jsonl":
            # JSONL datasets don't use split parameter
            return dataset_class(config)
        else:
            # HuggingFace datasets use split parameter
            return dataset_class(config, split=split)
    
    @staticmethod
    def register_dataset(mode: TrainingModeEnum, backend: str, dataset_class):
        """
        Register a custom dataset class.
        
        Args:
            mode: Training mode (e.g., TrainingModeEnum.PRETRAIN)
            backend: Backend name (e.g., "jsonl", "huggingface")
            dataset_class: Dataset class to register
        """
        DatasetFactory.DATASETS[(mode, backend)] = dataset_class
    
    @staticmethod
    def list_available_datasets():
        """List all available dataset configurations."""
        configs = {}
        for (mode, backend), dataset_class in DatasetFactory.DATASETS.items():
            key = f"{mode.value}_{backend}"
            configs[key] = {
                "mode": mode.value,
                "backend": backend,
                "class": dataset_class.__name__,
            }
        return configs