"""
Spark SQL type system — mirrors pyspark.sql.types.
Includes Polars <-> Spark type converters.
"""
from __future__ import annotations

import polars as pl


# ---------------------------------------------------------------------------
# Base + atomic types
# ---------------------------------------------------------------------------

class DataType:
    def simpleString(self) -> str:
        return type(self).__name__.replace("Type", "").lower()

    def __repr__(self) -> str:
        return f"{type(self).__name__}()"

    def __eq__(self, other):
        return type(self) == type(other)

    def __hash__(self):
        return hash(type(self).__name__)


class NullType(DataType):
    def simpleString(self): return "void"

class StringType(DataType):
    def simpleString(self): return "string"

class BinaryType(DataType):
    def simpleString(self): return "binary"

class BooleanType(DataType):
    def simpleString(self): return "boolean"

class ByteType(DataType):
    def simpleString(self): return "tinyint"

class ShortType(DataType):
    def simpleString(self): return "smallint"

class IntegerType(DataType):
    def simpleString(self): return "int"

class LongType(DataType):
    def simpleString(self): return "bigint"

class FloatType(DataType):
    def simpleString(self): return "float"

class DoubleType(DataType):
    def simpleString(self): return "double"

class DateType(DataType):
    def simpleString(self): return "date"

class TimestampType(DataType):
    def simpleString(self): return "timestamp"

class TimestampNTZType(DataType):
    def simpleString(self): return "timestamp_ntz"

class DecimalType(DataType):
    def __init__(self, precision: int = 10, scale: int = 0):
        self.precision = precision
        self.scale = scale

    def simpleString(self):
        return f"decimal({self.precision},{self.scale})"

    def __eq__(self, other):
        return isinstance(other, DecimalType) and self.precision == other.precision and self.scale == other.scale

    def __hash__(self):
        return hash(("DecimalType", self.precision, self.scale))


# ---------------------------------------------------------------------------
# Complex types
# ---------------------------------------------------------------------------

class ArrayType(DataType):
    def __init__(self, elementType: DataType, containsNull: bool = True):
        self.elementType = elementType
        self.containsNull = containsNull

    def simpleString(self):
        return f"array<{self.elementType.simpleString()}>"

    def __repr__(self):
        return f"ArrayType({self.elementType!r}, {self.containsNull})"

    def __eq__(self, other):
        return isinstance(other, ArrayType) and self.elementType == other.elementType

    def __hash__(self):
        return hash(("ArrayType", self.elementType))


class MapType(DataType):
    def __init__(self, keyType: DataType, valueType: DataType, valueContainsNull: bool = True):
        self.keyType = keyType
        self.valueType = valueType
        self.valueContainsNull = valueContainsNull

    def simpleString(self):
        return f"map<{self.keyType.simpleString()},{self.valueType.simpleString()}>"

    def __repr__(self):
        return f"MapType({self.keyType!r}, {self.valueType!r}, {self.valueContainsNull})"

    def __eq__(self, other):
        return (isinstance(other, MapType) and self.keyType == other.keyType
                and self.valueType == other.valueType)

    def __hash__(self):
        return hash(("MapType", self.keyType, self.valueType))


class StructField:
    def __init__(self, name: str, dataType: DataType, nullable: bool = True, metadata: dict | None = None):
        self.name = name
        self.dataType = dataType
        self.nullable = nullable
        self.metadata = metadata or {}

    def simpleString(self):
        return f"{self.name}:{self.dataType.simpleString()}"

    def __repr__(self):
        return f"StructField('{self.name}', {self.dataType!r}, {self.nullable})"

    def __eq__(self, other):
        return (isinstance(other, StructField) and self.name == other.name
                and self.dataType == other.dataType)

    def __hash__(self):
        return hash((self.name, self.dataType))


class StructType(DataType):
    def __init__(self, fields: list[StructField] | None = None):
        self.fields = fields or []

    def add(self, name: str, dataType: DataType, nullable: bool = True) -> "StructType":
        if isinstance(dataType, str):
            dataType = _parse_type_string(dataType)
        self.fields.append(StructField(name, dataType, nullable))
        return self

    def fieldNames(self) -> list[str]:
        return [f.name for f in self.fields]

    def simpleString(self):
        inner = ",".join(f.simpleString() for f in self.fields)
        return f"struct<{inner}>"

    def __repr__(self):
        return f"StructType([{', '.join(repr(f) for f in self.fields)}])"

    def __iter__(self):
        return iter(self.fields)

    def __len__(self):
        return len(self.fields)

    def __getitem__(self, key):
        if isinstance(key, str):
            for f in self.fields:
                if f.name == key:
                    return f
            raise KeyError(key)
        return self.fields[key]

    def __eq__(self, other):
        return isinstance(other, StructType) and self.fields == other.fields

    def __hash__(self):
        return hash(tuple(self.fields))


# ---------------------------------------------------------------------------
# Polars <-> Spark type converters
# ---------------------------------------------------------------------------

_SPARK_TO_POLARS: dict[type, pl.PolarsDataType] = {
    NullType:      pl.Null,
    StringType:    pl.Utf8,
    BinaryType:    pl.Binary,
    BooleanType:   pl.Boolean,
    ByteType:      pl.Int8,
    ShortType:     pl.Int16,
    IntegerType:   pl.Int32,
    LongType:      pl.Int64,
    FloatType:     pl.Float32,
    DoubleType:    pl.Float64,
    DateType:      pl.Date,
    TimestampType: pl.Datetime,
}

_POLARS_TO_SPARK: dict[pl.PolarsDataType, DataType] = {
    pl.Null:         NullType(),
    pl.Utf8:         StringType(),
    pl.String:       StringType(),
    pl.Binary:       BinaryType(),
    pl.Boolean:      BooleanType(),
    pl.Int8:         ByteType(),
    pl.Int16:        ShortType(),
    pl.Int32:        IntegerType(),
    pl.Int64:        LongType(),
    pl.UInt8:        ShortType(),
    pl.UInt16:       IntegerType(),
    pl.UInt32:       LongType(),
    pl.UInt64:       LongType(),
    pl.Float32:      FloatType(),
    pl.Float64:      DoubleType(),
    pl.Date:         DateType(),
    pl.Datetime:     TimestampType(),
}


def spark_type_to_polars(spark_type: DataType) -> pl.PolarsDataType:
    """Convert a Spark DataType to a Polars dtype."""
    if isinstance(spark_type, DecimalType):
        return pl.Decimal(spark_type.precision, spark_type.scale)
    if isinstance(spark_type, ArrayType):
        inner = spark_type_to_polars(spark_type.elementType)
        return pl.List(inner)
    if isinstance(spark_type, StructType):
        return pl.Struct([
            pl.Field(f.name, spark_type_to_polars(f.dataType))
            for f in spark_type.fields
        ])
    return _SPARK_TO_POLARS.get(type(spark_type), pl.Utf8)


def polars_dtype_to_spark(dtype: pl.PolarsDataType) -> DataType:
    """Convert a Polars dtype to a Spark DataType."""
    # Handle parameterized types
    if isinstance(dtype, pl.List):
        return ArrayType(polars_dtype_to_spark(dtype.inner))
    if isinstance(dtype, pl.Struct):
        fields = [StructField(f.name, polars_dtype_to_spark(f.dtype)) for f in dtype.fields]
        return StructType(fields)
    if isinstance(dtype, pl.Decimal):
        return DecimalType(dtype.precision or 38, dtype.scale or 0)
    if isinstance(dtype, pl.Datetime):
        return TimestampType()
    # Lookup base type
    for polars_t, spark_t in _POLARS_TO_SPARK.items():
        if dtype == polars_t:
            return spark_t
    return StringType()  # fallback


def polars_dtype_to_spark_str(dtype: pl.PolarsDataType) -> str:
    return polars_dtype_to_spark(dtype).simpleString()


def polars_schema_to_spark(schema: dict[str, pl.PolarsDataType]) -> StructType:
    fields = [StructField(name, polars_dtype_to_spark(dtype)) for name, dtype in schema.items()]
    return StructType(fields)


def _parse_type_string(s: str) -> DataType:
    """Parse simple type strings like 'string', 'int', 'bigint', etc."""
    mapping = {
        "string": StringType(), "str": StringType(),
        "int": IntegerType(), "integer": IntegerType(),
        "long": LongType(), "bigint": LongType(),
        "short": ShortType(), "smallint": ShortType(),
        "byte": ByteType(), "tinyint": ByteType(),
        "float": FloatType(),
        "double": DoubleType(),
        "boolean": BooleanType(), "bool": BooleanType(),
        "date": DateType(),
        "timestamp": TimestampType(),
        "binary": BinaryType(),
        "void": NullType(), "null": NullType(),
    }
    return mapping.get(s.strip().lower(), StringType())
