from dataclasses import dataclass, field
from typing import Any, Dict, Optional

from qwen_vl_utils.vision_process import IMAGE_FACTOR, MAX_PIXELS, MIN_PIXELS


@dataclass
class LerobotConfig:
    """Configuration for the LeRobot preprocessing pipeline.

    Dataset-specific camera display names are optional config inputs. Other
    dataset behavior is derived from the current LeRobot sample.
    """

    # Image resolution settings for different views
    resolution: Dict[str, int] = field(
        default_factory=lambda: {
            "face_view": -1,
            "left_wrist_view": 128,
            "right_wrist_view": 128,
        }
    )

    # Dataset splitting
    train_test_split: float = 0.9
    seed: int = 42

    # Instruction handling
    priority_order: Optional[Dict[str, float]] = None
    camera_name_mapping: Optional[Dict[str, str]] = None

    # Vision model parameters
    model_type: str = "qwen2_5"
    max_pixels: int = MAX_PIXELS
    min_pixels: int = MIN_PIXELS
    image_factor: int = IMAGE_FACTOR

    generate_subtask_ratio: float = 0.0

    def __post_init__(self):
        """Post-initialization validation and setup."""
        # Validate train/test split
        if not 0 < self.train_test_split < 1:
            raise ValueError(
                f"train_test_split must be between 0 and 1, got {self.train_test_split}"
            )

    def as_dict(self) -> Dict:
        """Convert configuration to dictionary format.

        Returns:
            Dict: Configuration as dictionary
        """
        return self.__dict__

    def update(self, **kwargs) -> "LerobotConfig":
        """Update configuration parameters.

        Args:
            **kwargs: Key-value pairs to update

        Returns:
            LerobotConfig: Updated configuration instance
        """
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                raise ValueError(f"Unknown configuration parameter: {key}")
        return self

    def __getitem__(self, key: str):
        return getattr(self, key)

    @classmethod
    def from_yaml_dict(cls, yaml_dict: Dict[str, Any]) -> "LerobotConfig":
        """
        Build a LerobotConfig instance from a YAML dictionary.

        Supports two styles:

        1) Top-level fields:
            train_test_split: 0.8
            model_type: qwen2_5

        2) Nested under `data:` (higher priority):
            data:
              train_test_split: 0.8
              model_type: qwen2_5

        Keys inside `data:` override top-level keys.
        """

        data_config = yaml_dict.get("data", {})

        def get(key: str, default: Any = None):
            """
            Helper function:
            Read from `data` first, then fallback to the top-level YAML.
            """
            return data_config.get(key, yaml_dict.get(key, default))

        # Construct only fields that actually exist in LerobotConfig
        params: Dict[str, Any] = {
            # Action prediction settings
            # Image resolution per camera view
            "resolution": get(
                "resolution",
                {
                    "face_view": -1,
                    "left_wrist_view": 128,
                    "right_wrist_view": 128,
                },
            ),
            # Dataset train/test split configuration
            "train_test_split": get("train_test_split", 0.9),
            "seed": get("seed", 42),
            # Instruction priority ordering (optional)
            "priority_order": get("priority_order", None),
            "camera_name_mapping": get("camera_name_mapping", None),
            # Vision model parameters
            "model_type": get("model_type", "qwen2_5"),
            "max_pixels": get("max_pixels", MAX_PIXELS),
            "min_pixels": get("min_pixels", MIN_PIXELS),
            "image_factor": get("image_factor", IMAGE_FACTOR),
            # Subtask generation ratio
            "generate_subtask_ratio": get("generate_subtask_ratio", 0.0),
        }

        # Keep only valid dataclass fields (ignore unknown YAML keys)
        valid_fields = {f.name for f in cls.__dataclass_fields__.values()}
        filtered_params = {k: v for k, v in params.items() if k in valid_fields}

        return cls(**filtered_params)