"""
DataFrame factory service for SparkSession.

This service handles DataFrame creation, schema inference, and validation
following the Single Responsibility Principle.
"""

from typing import Any, Dict, List, Optional, Tuple, Union
from mock_spark.spark_types import (
    StructType,
    StructField,
    StringType,
)
from mock_spark.dataframe import DataFrame
from mock_spark.session.config import SparkConfig
from mock_spark.core.exceptions import IllegalArgumentException


class DataFrameFactory:
    """Factory for creating DataFrames with validation and coercion."""

    def create_dataframe(
        self,
        data: Union[List[Dict[str, Any]], List[Any]],
        schema: Optional[Union[StructType, List[str], str]],
        engine_config: SparkConfig,
        storage: Any,
    ) -> DataFrame:
        """Create a DataFrame from data.

        Args:
            data: List of dictionaries or tuples representing rows.
            schema: Optional schema definition (StructType or list of column names).
            engine_config: Engine configuration for validation and coercion.
            storage: Storage manager for the DataFrame.

        Returns:
            DataFrame instance with the specified data and schema.

        Raises:
            IllegalArgumentException: If data is not in the expected format.

        Example:
            >>> data = [{"name": "Alice", "age": 25}, {"name": "Bob", "age": 30}]
            >>> df = factory.create_dataframe(data, None, config, storage)
        """
        if not isinstance(data, list):
            raise IllegalArgumentException(
                "Data must be a list of dictionaries or tuples"
            )

        # Handle PySpark StructType - convert to StructType
        # Check if it's a PySpark StructType (has 'fields' attribute but not StructType)
        if (
            schema is not None
            and not isinstance(schema, (StructType, str, list))
            and hasattr(schema, "fields")
        ):
            # This is likely a PySpark StructType - convert it
            schema = self._convert_pyspark_struct_type(schema)

        # Handle DDL schema strings
        if isinstance(schema, str):
            from mock_spark.core.ddl_adapter import parse_ddl_schema

            schema = parse_ddl_schema(schema)

        # Handle list of column names as schema
        if isinstance(schema, list):
            # Convert tuples to dictionaries using provided column names first
            if data and isinstance(data[0], tuple):
                reordered_data = []
                column_names = schema
                for row in data:
                    if isinstance(row, tuple):
                        row_dict = {column_names[i]: row[i] for i in range(len(row))}
                        reordered_data.append(row_dict)
                    else:
                        reordered_data.append(row)
                data = reordered_data

                # Now infer schema from the converted data
                from mock_spark.core.schema_inference import SchemaInferenceEngine

                schema, data = SchemaInferenceEngine.infer_from_data(data)
            else:
                # For non-tuple data with column names, use StringType as default
                fields = [StructField(name, StringType()) for name in schema]
                schema = StructType(fields)

        if schema is None:
            # Infer schema from data using SchemaInferenceEngine
            if not data:
                # For empty dataset, create empty schema
                schema = StructType([])
            else:
                # Check if data is in expected format
                sample_row = data[0]
                if not isinstance(sample_row, (dict, tuple)):
                    raise IllegalArgumentException(
                        "Data must be a list of dictionaries or tuples"
                    )

                if isinstance(sample_row, dict):
                    # Use SchemaInferenceEngine for dictionary data
                    from mock_spark.core.schema_inference import SchemaInferenceEngine

                    schema, data = SchemaInferenceEngine.infer_from_data(data)
                elif isinstance(sample_row, tuple):
                    # For tuples, we need column names - this should have been handled earlier
                    # If we get here, it's an error
                    raise IllegalArgumentException(
                        "Cannot infer schema from tuples without column names. "
                        "Please provide schema or use list of column names."
                    )

        # Apply validation and optional type coercion per mode
        # Note: When an explicit schema is provided with empty data, we still need to preserve
        # the schema even though validation is skipped (since there's no data to validate)
        if isinstance(schema, StructType) and data:
            from mock_spark.core.data_validation import DataValidator

            validator = DataValidator(
                schema,
                validation_mode=engine_config.validation_mode,
                enable_coercion=engine_config.enable_type_coercion,
            )

            # Validate if in strict mode
            if engine_config.validation_mode == "strict":
                validator.validate(data)

            # Coerce if enabled
            if engine_config.enable_type_coercion:
                data = validator.coerce(data)

        # Ensure schema is always StructType at this point
        # IMPORTANT: When explicit schema is provided with empty data, preserve it!
        if not isinstance(schema, StructType):
            # This should never happen, but provide a fallback
            schema = StructType([])

        # Validate that schema is properly initialized with fields attribute
        # This ensures empty DataFrames with explicit schemas preserve column information
        if isinstance(schema, StructType) and not hasattr(schema, "fields"):
            # This shouldn't happen, but handle edge case
            schema = StructType([])
            # fields can be empty list, but that's valid for empty schemas
            # If schema was provided explicitly, trust it even if fields is empty

        return DataFrame(data, schema, storage)

    def _handle_schema_inference(
        self, data: List[Dict[str, Any]], schema: Optional[Any]
    ) -> Tuple[StructType, List[Dict[str, Any]]]:
        """Handle schema inference or conversion.

        Args:
            data: List of dictionaries representing rows.
            schema: Optional schema definition.

        Returns:
            Tuple of (inferred_schema, normalized_data).
        """
        if schema is None:
            from mock_spark.core.schema_inference import SchemaInferenceEngine

            return SchemaInferenceEngine.infer_from_data(data)
        else:
            # Schema provided, return as-is
            return schema, data

    def _apply_validation_and_coercion(
        self,
        data: List[Dict[str, Any]],
        schema: StructType,
        engine_config: SparkConfig,
    ) -> List[Dict[str, Any]]:
        """Apply validation and type coercion.

        Args:
            data: List of dictionaries representing rows.
            schema: Schema to validate against.
            engine_config: Engine configuration.

        Returns:
            Validated and coerced data.
        """
        from mock_spark.core.data_validation import DataValidator

        validator = DataValidator(
            schema,
            validation_mode=engine_config.validation_mode,
            enable_coercion=engine_config.enable_type_coercion,
        )

        # Validate if in strict mode
        if engine_config.validation_mode == "strict":
            validator.validate(data)

        # Coerce if enabled
        if engine_config.enable_type_coercion:
            data = validator.coerce(data)

        return data

    def _convert_pyspark_struct_type(self, pyspark_schema: Any) -> StructType:
        """Convert PySpark StructType to StructType.

        Args:
            pyspark_schema: PySpark StructType object with 'fields' attribute.

        Returns:
            StructType equivalent.
        """
        from mock_spark.spark_types import (
            DataType,
            StringType,
            IntegerType,
            LongType,
            FloatType,
            DoubleType,
            BooleanType,
            TimestampType,
            DateType,
            DecimalType,
            ArrayType,
            MapType,
            StructType,
        )

        def convert_pyspark_field(field: Any) -> StructField:
            """Convert PySpark StructField to StructField."""
            field_name = field.name
            field_nullable = getattr(field, "nullable", True)

            # Convert PySpark data type to MockSpark data type
            pyspark_type = field.dataType
            mock_type = convert_pyspark_data_type(pyspark_type)

            return StructField(
                name=field_name, dataType=mock_type, nullable=field_nullable
            )

        def convert_pyspark_data_type(pyspark_type: Any) -> DataType:
            """Convert PySpark DataType to DataType."""
            # Get the type name as string for comparison
            type_name = type(pyspark_type).__name__

            if type_name == "StringType":
                return StringType()
            elif type_name == "IntegerType":
                return IntegerType()
            elif type_name == "LongType":
                return LongType()
            elif type_name == "FloatType":
                return FloatType()
            elif type_name == "DoubleType":
                return DoubleType()
            elif type_name == "BooleanType":
                return BooleanType()
            elif type_name == "TimestampType":
                return TimestampType()
            elif type_name == "DateType":
                return DateType()
            elif type_name == "DecimalType":
                precision = getattr(pyspark_type, "precision", 10)
                scale = getattr(pyspark_type, "scale", 0)
                return DecimalType(precision=precision, scale=scale)
            elif type_name == "ArrayType":
                element_type = convert_pyspark_data_type(pyspark_type.elementType)
                return ArrayType(element_type)
            elif type_name == "MapType":
                key_type = convert_pyspark_data_type(pyspark_type.keyType)
                value_type = convert_pyspark_data_type(pyspark_type.valueType)
                return MapType(key_type, value_type)
            elif type_name == "StructType":
                # Recursive conversion for nested structs
                fields = [convert_pyspark_field(f) for f in pyspark_type.fields]
                return StructType(fields)
            else:
                # Default to StringType for unknown types
                return StringType()

        # Convert all fields
        fields = [convert_pyspark_field(field) for field in pyspark_schema.fields]
        return StructType(fields)
