"""
DataFrame factory service for SparkSession.

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

from typing import Any, Optional, 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)
