from typing import Dict, Any, Optional, Union, Literal, List
from pydantic import BaseModel, Field, validator, field_validator


class DatasetOutputs(BaseModel):
    rebalanced: Optional[str] = None
    train_split: Optional[str] = None
    test_split: Optional[str] = None


class ProcessedOutputConfig(BaseModel):
    storage_ref: str
    base_path: str
    train_dataset: str
    test_dataset: Optional[str] = None


class PreprocessingConfig(BaseModel):
    downsample: Optional[int] = None
    rebalance: Optional[Dict[str, Union[int, float]]] = None
    split_ratio: Optional[int] = None
    output_config: Optional[ProcessedOutputConfig] = None
    # shuffle: Optional[bool] = True
    # random_seed: Optional[int] = 42

    @field_validator('rebalance')
    @classmethod
    def validate_rebalance(cls, v):
        if v is None:
            return v

        if not isinstance(v, dict) or not v:
            raise ValueError('Rebalance must be a non-empty dictionary')

        values = list(v.values())

        # Check if all values are integers (absolute mode)
        if all(isinstance(val, int) for val in values):
            if any(val < 0 for val in values):
                raise ValueError("Absolute rebalance values must be non-negative integers")
            return v

        # Check if all values are floats/numbers between 0 and 1 (fraction mode)
        if all(isinstance(val, (int, float)) for val in values):
            if any(not (0 <= val <= 1) for val in values):
                raise ValueError("Fraction rebalance values must be between 0 and 1")

            total = sum(values)
            if not abs(total - 1.0) < 1e-6:
                raise ValueError(f"Fraction rebalance values must sum to 1.0, got {total}")
            return v
            # Mixed types not allowed

        raise ValueError("All rebalance values must be either integers (absolute) or floats (fraction)")


class DatasetConfig(BaseModel):
    storage_ref: str = Field(description="Reference to storage configuration")
    storage_path: str
    preprocessing: Optional[PreprocessingConfig] = None


class DatasetConfigs(BaseModel):
    train_dataset: DatasetConfig
    test_dataset: Optional[DatasetConfig] = None