from __future__ import annotations

import base64
import collections.abc
import dataclasses
import decimal
import enum
from datetime import date, datetime, time, timedelta
from typing import Any, FrozenSet, List, Set, Tuple, Type, TypeVar, Union, cast

import attrs
import cattrs
import dateutil.parser
import isodate
from pydantic import BaseModel
from typing_extensions import get_args, get_origin, get_type_hints, is_typeddict

from chalk.features._encoding.primitive import TPrimitive
from chalk.utils.collections import is_namedtuple, unwrap_optional_if_needed
from chalk.utils.enum import get_enum_value_type

try:
    import numpy as np
except ImportError:
    np = None

TRich = TypeVar("TRich")
T = TypeVar("T")

# Converter to go from TPython to a TPrimitive
_rich_converter = cattrs.Converter()


#########
# Structs
#########

_rich_converter.register_unstructure_hook_func(
    is_namedtuple, lambda x: {k: _rich_converter.unstructure(v) for (k, v) in x._asdict().items()}
)
_rich_converter.register_unstructure_hook_func(
    dataclasses.is_dataclass,
    lambda x: {field.name: _rich_converter.unstructure(getattr(x, field.name)) for field in dataclasses.fields(x)},
)
_rich_converter.register_unstructure_hook(
    BaseModel, lambda x: {k: _rich_converter.unstructure(v) for (k, v) in x.dict().items()}
)
_rich_converter.register_unstructure_hook_func(
    attrs.has,
    lambda x: {k: _rich_converter.unstructure(v) for (k, v) in attrs.asdict(x, recurse=False).items()},
)


def _is_struct(typ: Type):
    typ = unwrap_optional_if_needed(typ)
    return (
        is_namedtuple(typ)
        or dataclasses.is_dataclass(typ)
        or is_typeddict(typ)
        or (isinstance(typ, type) and issubclass(typ, BaseModel))
        or attrs.has(typ)
    )


def _structure_struct(obj: Any, typ: Type):
    if not isinstance(obj, (collections.abc.Sequence, collections.abc.Mapping)):
        # If an already-constructed class is passed in, then we still need to restructure it
        # to ensure that any recursive fields are properly structured
        obj = _rich_converter.unstructure(obj)
    typ = unwrap_optional_if_needed(typ)
    type_hints = get_type_hints(typ)
    if obj is None:
        obj = {k: None for k in type_hints}
    if isinstance(obj, collections.abc.Sequence):
        obj = {k: v for (k, v) in zip(type_hints.keys(), obj)}
    if not isinstance(obj, collections.abc.Mapping):
        raise TypeError(f"Unable to structure object `{obj}` into type `{typ}`. Only dictionaries are supported.")
    kwargs = {k: _rich_converter.structure(v, type_hints[k]) for (k, v) in obj.items()}
    if isinstance(typ, type) and issubclass(typ, BaseModel):
        # Using .construct to bypass pydantic validation, which could fail if there are null values in kwargs
        # and the field is not annotated as null
        return typ.construct(**kwargs)
    else:
        return typ(**kwargs)


_rich_converter.register_structure_hook_func(_is_struct, _structure_struct)

#############
# Collections
#############

# Unlike lists and tuples which are ordered, sets are unordered
# Thus, when converting to a list, the order must be consistent
# To do so, sort all the values
_rich_converter.register_unstructure_hook(set, lambda x: sorted(list(x)))
_rich_converter.register_unstructure_hook(frozenset, lambda x: sorted(list(x)))

# Lists and tuples do not need an unstructure hook -- the default is fine


def _structure_collection(
    obj: Any, typ: Union[Type[FrozenSet[T]], Type[Set[T]], Type[List[T]], Type[Tuple[T, ...]]]
) -> Union[FrozenSet[T], Set[T], List[T], Tuple[T, ...]]:
    origin = get_origin(typ)
    if origin in (set, Set):
        constructor = set
    elif origin in (frozenset, FrozenSet):
        constructor = frozenset
    elif origin in (list, List):
        constructor = list
    elif origin in (tuple, Tuple):
        constructor = tuple
    else:
        raise TypeError(f"Unsupported set type: {typ}")
    args = get_args(typ)
    if len(args) < 1:
        raise TypeError(
            f"{typ} types must be parameterized with the type of the contained value -- for example, `{typ}[int]`"
        )
    if len(args) > 1:
        if origin in (tuple, Tuple):
            if len(args) != 2 and args[1] != ...:
                raise TypeError(
                    (
                        "Tuple features must have a fixed type and be variable-length tuples (e.g. `Tuple[int, ...]`). "
                        " If you would like a fixed-length of potentially different types, used a NamedTuple."
                    )
                )
        else:
            raise TypeError(f"{typ} should be parameterized with only one type")

    inner_typ = args[0]
    if obj is None:
        return cast(Union[FrozenSet[T], Set[T], List[T], Tuple[T, ...]], None)
    if not isinstance(obj, (collections.abc.Set, collections.abc.Sequence)):
        raise TypeError(f"Cannot structure '{obj}' into a `{typ}`")
    return constructor(_rich_converter.structure(x, inner_typ) for x in obj)


def _is_collection(typ: Type):
    origin = get_origin(typ)
    return origin in (set, Set, frozenset, FrozenSet, list, List, tuple, Tuple)


_rich_converter.register_structure_hook_func(_is_collection, _structure_collection)

#######
# Enums
#######

_rich_converter.register_unstructure_hook(enum.Enum, lambda x: _rich_converter.unstructure(x.value))


def _structure_enum(obj: Any, typ: Type[enum.Enum]) -> enum.Enum:
    if isinstance(obj, typ):
        return obj
    if obj is None:
        return cast(enum.Enum, None)
    enum_typ = get_enum_value_type(typ)
    try:
        return typ(_rich_converter.structure(obj, enum_typ))
    except (TypeError, ValueError):
        pass
    if isinstance(obj, str):
        try:
            return typ[obj]
        except KeyError:
            pass
    allowed_values = ", ".join(f"'{x}'" for x in typ.__members__.values())
    raise ValueError(f"Cannot convert '{obj}' to Enum `{typ}`. Possible values are: {allowed_values}")


_rich_converter.register_structure_hook(enum.Enum, _structure_enum)

##########
# Decimals
##########

_rich_converter.register_unstructure_hook(decimal.Decimal, lambda x: str(x.normalize()))


def _structure_decimal(obj: Any, typ: Type[decimal.Decimal]) -> decimal.Decimal:
    if obj is None:
        return cast(decimal.Decimal, None)
    if isinstance(obj, decimal.Decimal):
        return obj
    return decimal.Decimal(obj)


_rich_converter.register_structure_hook(decimal.Decimal, _structure_decimal)


#######
# Bytes
#######

_rich_converter.register_unstructure_hook(bytes, lambda x: x)


def _structure_bytes(x: Any, typ: Type[bytes]) -> bytes:
    if x is None:
        return cast(bytes, None)
    if isinstance(x, bytes):
        return x
    if isinstance(x, str):
        return base64.b64decode(x)
    raise TypeError(f"Cannot structure {x} into bytes")


_rich_converter.register_structure_hook(bytes, _structure_bytes)

##########
# Duration
##########

_rich_converter.register_unstructure_hook(timedelta, lambda x: x)


def _structure_timedelta(x: Any, typ: Type[timedelta]) -> timedelta:
    if x is None:
        return cast(timedelta, None)
    if isinstance(x, timedelta):
        return x
    if isinstance(x, str):
        return isodate.parse_duration(x)
    raise TypeError(f"Cannot structure {x} into a duration")


_rich_converter.register_structure_hook(timedelta, _structure_timedelta)


######
# Date
######

_rich_converter.register_unstructure_hook_func(
    lambda x: isinstance(x, date) and not isinstance(x, datetime), lambda x: datetime.combine(x, time())
)


def _structure_date(x: Any, typ: Type[date]) -> date:
    if x is None:
        return cast(date, None)
    if isinstance(x, datetime):
        if x.time() != time():
            raise TypeError(f"Datetime '{x}' has a non-zero time component, which cannot be safely cast into a date")
        return x.date()
    if isinstance(x, date):
        return x
    if isinstance(x, str):
        return isodate.parse_date(x)
    raise TypeError(f"Cannot structure {x} into a date")


_rich_converter.register_structure_hook_func(
    lambda x: isinstance(x, type) and issubclass(x, date) and not issubclass(x, datetime), _structure_date
)

######
# Time
######

_rich_converter.register_unstructure_hook(time, lambda x: x)


def _structure_time(x: Any, typ: Type[time]) -> time:
    if x is None:
        return cast(time, None)
    if isinstance(x, time):
        return x
    if isinstance(x, str):
        return isodate.parse_time(x)
    raise TypeError(f"Cannot structure {x} into a time")


_rich_converter.register_structure_hook(time, _structure_time)


##########
# Datetime
##########

_rich_converter.register_unstructure_hook(datetime, lambda x: x)


def _structure_datetime(x: Any, typ: Type[datetime]) -> datetime:
    if x is None:
        return cast(datetime, None)
    if isinstance(x, str):
        x = dateutil.parser.parse(x)
    if isinstance(x, datetime):
        return x
    if isinstance(x, date):
        return datetime.combine(x, time())
    raise TypeError(f"Cannot structure {x} into a datetime")


_rich_converter.register_structure_hook(datetime, _structure_datetime)


#############
# Numpy types
#############
_rich_converter.register_unstructure_hook_func(lambda x: np is not None and np.issubdtype(x, np.bool_), bool)
_rich_converter.register_unstructure_hook_func(lambda x: np is not None and np.issubdtype(x, np.integer), int)
_rich_converter.register_unstructure_hook_func(lambda x: np is not None and np.issubdtype(x, np.floating), float)


##############
# Rich types
##############


def _structure_rich(obj: Any, typ: Type):
    if obj is None:
        # Always allow None, even if the field is non-optional
        return None
    if issubclass(typ, enum.Enum):
        return _structure_enum(obj, typ)
    if isinstance(obj, typ):
        return obj
    if issubclass(typ, bool):
        # For booleans, we don't want to use the bool(...) constructor, as
        # that doesn't handle strings like y/yes/n/no/t/true/f/false
        if isinstance(obj, str):
            if obj.lower() in ("y", "yes", "t", "true", "1"):
                return True
            if obj.lower() in ("n", "no", "f", "false", "0"):
                return False
        if isinstance(obj, int):
            if obj == 0:
                return False
            if obj == 1:
                return True
        raise ValueError(
            f"Cannot convert {obj} to a boolean. Allowed values are y, yes, t, true, 1, n, no, f, false, 0."
        )

    return typ(obj)


_rich_converter.register_structure_hook(int, _structure_rich)
_rich_converter.register_structure_hook(float, _structure_rich)
_rich_converter.register_structure_hook(bytes, _structure_rich)
_rich_converter.register_structure_hook(str, _structure_rich)
_rich_converter.register_structure_hook(bool, _structure_rich)


def unstructure_rich_to_primitive(val: Any) -> TPrimitive:
    return _rich_converter.unstructure(val)


def structure_primitive_to_rich(val: TPrimitive, typ: Type[TRich]) -> TRich:
    return _rich_converter.structure(val, typ)
