from __future__ import annotations

import enum
import json
from typing import List, Literal, Optional, Union

import pyarrow as pa
from pydantic import BaseConfig, BaseModel, Field
from typing_extensions import Annotated

__all__ = [
    "deserialize_dtype",
    "serialize_dtype",
    "serialize_pyarrow_dtype",
]


class _DTypeCode(enum.IntEnum):
    # NULL = 1
    BOOL = 2
    INT8 = 3
    INT16 = 4
    INT32 = 5
    INT64 = 6
    UINT8 = 7
    UINT16 = 8
    UINT32 = 9
    UINT64 = 10
    FLOAT16 = 11
    FLOAT32 = 12
    FLOAT64 = 13
    TIME32 = 14  # time unit
    TIME64 = 15  # time unit
    TIMESTAMP = 16  # time unit, time zone
    DATE32 = 17
    DATE64 = 18
    DURATION = 19  # time unit
    BINARY = 20  # length
    STRING = 21
    LARGE_BINARY = 22
    LARGE_STRING = 23
    # Decimal is not yet supported by chalk
    # DECIMAL128 = 24  # precision, scale
    LIST = 25  # value type, size
    LARGE_LIST = 26  # value type
    STRUCT = 27  # fields


class _TimeUnitCode(str, enum.Enum):
    SECOND = "s"
    MILLISECONDS = "ms"
    MICROSECONDS = "us"
    NANOSECONDS = "ns"


class _BaseSerializedDType(BaseModel):
    type_code: _DTypeCode

    def to_pyarrow_dtype(self) -> pa.DataType:
        raise NotImplementedError

    class Config(BaseConfig):
        use_enum_values = True


class _SingletonDType(_BaseSerializedDType):
    type_code: Literal[
        _DTypeCode.BOOL,
        _DTypeCode.INT8,
        _DTypeCode.INT16,
        _DTypeCode.INT32,
        _DTypeCode.INT64,
        _DTypeCode.UINT8,
        _DTypeCode.UINT16,
        _DTypeCode.UINT32,
        _DTypeCode.UINT64,
        _DTypeCode.FLOAT16,
        _DTypeCode.FLOAT32,
        _DTypeCode.FLOAT64,
        _DTypeCode.DATE32,
        _DTypeCode.DATE64,
        _DTypeCode.STRING,
        _DTypeCode.LARGE_STRING,
        _DTypeCode.LARGE_BINARY,
    ]

    def to_pyarrow_dtype(self) -> pa.DataType:
        if self.type_code == _DTypeCode.BOOL:
            return pa.bool_()
        try:
            return getattr(pa, self.type_code.name.lower())()
        except AttributeError:
            raise ValueError(f"Unsupported dtype: {self.type_code}") from None


class _TimeUnitDType(_BaseSerializedDType):
    type_code: Literal[_DTypeCode.TIME32, _DTypeCode.TIME64, _DTypeCode.DURATION]
    time_unit: _TimeUnitCode

    def to_pyarrow_dtype(self) -> pa.DataType:
        return getattr(pa, self.type_code.name.lower())(self.time_unit)


class _TimestampDType(_BaseSerializedDType):
    type_code: Literal[_DTypeCode.TIMESTAMP]
    time_unit: _TimeUnitCode
    timezone: Optional[str]

    def to_pyarrow_dtype(self) -> pa.DataType:
        return pa.timestamp(self.time_unit, self.timezone)


class _BinaryDType(_BaseSerializedDType):
    type_code: Literal[_DTypeCode.BINARY]
    length: int

    def to_pyarrow_dtype(self) -> pa.DataType:
        return pa.binary(self.length)


class _ListDType(_BaseSerializedDType):
    type_code: Literal[_DTypeCode.LIST]
    inner_dtype: _SerializedDType
    length: int

    def to_pyarrow_dtype(self) -> pa.DataType:
        return pa.list_(self.inner_dtype.to_pyarrow_dtype(), self.length)


class _LargeListDType(_BaseSerializedDType):
    type_code: Literal[_DTypeCode.LARGE_LIST]
    inner_dtype: _SerializedDType

    def to_pyarrow_dtype(self) -> pa.DataType:
        return pa.large_list(self.inner_dtype.to_pyarrow_dtype())


class _StructField(BaseModel):
    name: str
    dtype: _SerializedDType
    nullable: bool = True

    def to_pyarrow_dtype(self) -> pa.Field:
        return pa.field(self.name, self.dtype.to_pyarrow_dtype(), self.nullable)


class _StructDType(_BaseSerializedDType):
    type_code: Literal[_DTypeCode.STRUCT]
    fields: List[_StructField]

    def to_pyarrow_dtype(self) -> pa.Field:
        return pa.struct([x.to_pyarrow_dtype() for x in self.fields])


_SerializedDType = Annotated[
    Union[_SingletonDType, _TimeUnitDType, _TimestampDType, _BinaryDType, _ListDType, _LargeListDType, _StructDType],
    Field(discriminator="type_code"),
]

_StructField.update_forward_refs()
_ListDType.update_forward_refs()
_LargeListDType.update_forward_refs()


class _Deserializer(BaseModel):
    dtype: _SerializedDType


def deserialize_dtype(serialized_dtype: str) -> pa.DataType:
    deserializer = _Deserializer(dtype=json.loads(serialized_dtype))
    pa_dtype = deserializer.dtype.to_pyarrow_dtype()
    return pa_dtype


def serialize_pyarrow_dtype(pa_dtype: pa.DataType) -> _SerializedDType:
    """Deprecated -- use `serialize_dtype` instead. For backwards compatibility with Engine"""
    return _serialize_pyarrow_dtype(pa_dtype)


def serialize_dtype(pa_dtype: pa.DataType) -> str:
    return _serialize_pyarrow_dtype(pa_dtype).json()


def _serialize_pyarrow_dtype(pa_dtype: pa.DataType) -> _SerializedDType:
    if pa.types.is_boolean(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.BOOL)
    if pa.types.is_uint8(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.UINT8)
    if pa.types.is_uint16(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.UINT16)
    if pa.types.is_uint32(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.UINT32)
    if pa.types.is_uint64(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.UINT64)
    if pa.types.is_int8(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.INT8)
    if pa.types.is_int16(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.INT16)
    if pa.types.is_int32(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.INT32)
    if pa.types.is_int64(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.INT64)
    if pa.types.is_float16(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.FLOAT16)
    if pa.types.is_float32(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.FLOAT32)
    if pa.types.is_float64(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.FLOAT64)
    if pa.types.is_time32(pa_dtype):
        return _TimeUnitDType(type_code=_DTypeCode.TIME32, time_unit=_TimeUnitCode(pa_dtype.unit))

    if pa.types.is_time64(pa_dtype):
        return _TimeUnitDType(type_code=_DTypeCode.TIME64, time_unit=_TimeUnitCode(pa_dtype.unit))

    if pa.types.is_timestamp(pa_dtype):
        return _TimestampDType(
            type_code=_DTypeCode.TIMESTAMP,
            time_unit=_TimeUnitCode(pa_dtype.unit),
            timezone=pa_dtype.tz,
        )

    if pa.types.is_date32(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.DATE32)
    if pa.types.is_date64(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.DATE64)
    if pa.types.is_duration(pa_dtype):
        return _TimeUnitDType(type_code=_DTypeCode.DURATION, time_unit=_TimeUnitCode(pa_dtype.unit))

    if pa.types.is_large_string(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.LARGE_STRING)
    if pa.types.is_large_binary(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.LARGE_BINARY)
    if pa.types.is_fixed_size_binary(pa_dtype):
        return _BinaryDType(type_code=_DTypeCode.BINARY, length=pa_dtype.byte_width)
    if pa.types.is_binary(pa_dtype):
        return _BinaryDType(type_code=_DTypeCode.BINARY, length=-1)
    if pa.types.is_string(pa_dtype):
        return _SingletonDType(type_code=_DTypeCode.STRING)

    if pa.types.is_large_list(pa_dtype):
        return _LargeListDType(
            type_code=_DTypeCode.LARGE_LIST, inner_dtype=_serialize_pyarrow_dtype(pa_dtype.value_type)
        )
    if pa.types.is_fixed_size_list(pa_dtype):
        return _ListDType(
            type_code=_DTypeCode.LIST,
            inner_dtype=_serialize_pyarrow_dtype(pa_dtype.value_type),
            length=pa_dtype.list_size,
        )

    if pa.types.is_list(pa_dtype):
        return _ListDType(
            type_code=_DTypeCode.LIST,
            inner_dtype=_serialize_pyarrow_dtype(pa_dtype.value_type),
            length=-1,
        )

    if pa.types.is_struct(pa_dtype):
        schema = pa.schema(pa_dtype)
        return _StructDType(
            type_code=_DTypeCode.STRUCT,
            fields=[
                _StructField(name=name, dtype=_serialize_pyarrow_dtype(dtype))
                for (name, dtype) in zip(schema.names, schema.types)
            ],
        )

    raise ValueError(f"Unsupported dtype: {pa_dtype}")
