import dateutil.parser
import pytz

from copy import copy
from typing import Union, NamedTuple
from collections import OrderedDict, defaultdict
from collections.abc import Sequence, Mapping, MappingView
from functools import lru_cache
from datetime import datetime

from marshmallow import Schema, fields, pre_load, post_dump, EXCLUDE, ValidationError
from marshmallow_jsonschema import JSONSchema

from spintop.utils.marshmallow_dataclass import class_schema
from spintop.utils.marshmallow_oneofschema import OneOfSchema

from spintop.utils import local_tz, isnan, dict_ops

_TYPE_TO_CLS_MAP = {}
_SCHEMA_TYPES = {}

TYPE_KEY = '_type'

_TYPE_CLS_CHILDREN = defaultdict(set)

_REGISTERED_TYPES = []

def register_type(cls, _type=None, schema_cls=None):
    if _type is None:
        _type = cls.__name__
    cls._type = _type
    
    if _type in _TYPE_TO_CLS_MAP and _TYPE_TO_CLS_MAP[_type] != cls:
        raise ValueError("Type %s is already registered to class %s" % (_type, _TYPE_TO_CLS_MAP[_type]))

    _TYPE_TO_CLS_MAP[_type] = cls

    _REGISTERED_TYPES.append((cls, schema_cls))

    for schema_key in _SCHEMA_TYPES:
        Schema = DataClassSchema(cls, schema_type=get_base_schema_cls(schema_key), schema_cls=schema_cls)

        # Building the schema now improves serialization performance by up to 800%
        _SCHEMA_TYPES[schema_key].type_schemas[_type] = Schema()

    for other_cls in _TYPE_CLS_CHILDREN:
        if other_cls in cls.__mro__:
            _TYPE_CLS_CHILDREN[other_cls].add(cls)
    _TYPE_CLS_CHILDREN[cls] = _TYPE_CLS_CHILDREN[cls]

def cls_subclasses(cls):
    """Returns a set containing the subclasses of cls 
    (if cls was at some point passed to register_type as argument)"""
    return _TYPE_CLS_CHILDREN[cls]

def cls_of_obj(obj):
    if isinstance(obj, Mapping):
        return cls_of_serialized(obj)
    else:
        return obj.__class__
    # try:
    #     _type = type_of(obj)
    #     cls = cls_of_type(_type)
    #     if cls is None:
    #         raise AttributeError(f'{cls} is not mapped to a type')
    #     return cls
    # except AttributeError:
    #     return cls_of_serialized(obj)

def cls_of_serialized(obj_serialized):
    return cls_of_type(serialized_type_of(obj_serialized))

def cls_of_type(_type):
    return _TYPE_TO_CLS_MAP.get(_type, None)

def type_of(cls):
    return cls._type

def serialized_type_of(obj_serialized):
    return obj_serialized[TYPE_KEY]

def type_dict_of(cls):
    return {TYPE_KEY: type_of(cls)}

def is_type_of(cls, base_cls):
    return type_of(cls) == type_of(base_cls)

def is_serialized_type_of(obj_serialized, base_cls):
    return serialized_type_of(obj_serialized) == type_of(base_cls)

def get_serializer(schema_key=None):
    return Serializer(get_base_schema_cls(schema_key))

@lru_cache(maxsize=None)
def DataClassSchema(cls, schema_type=None, schema_cls=None):
    if schema_cls:
        schema_cls = compose_schema_cls(schema_type, schema_cls)
    else:
        schema_cls = schema_type
    schema = class_schema(cls, base_schema=schema_cls)
    return schema

def get_base_schema_cls(schema_key):
    try:
        return _SCHEMA_TYPES[schema_key]
    except KeyError:
        raise ValueError(f'Schema key {schema_key!r} not in supported list: {_SCHEMA_TYPES.keys()!r}')

def compose_schema_cls(base, child):
    return type(f'Composed{child.__name__}', (base, child), {})

class Serializer(object):
    def __init__(self, base_schema_cls):
        self.base_schema_cls = base_schema_cls

    def schema(self, *args, **kwargs):
        return self.base_schema_cls(*args, **kwargs)

    def serialize(self, obj, cls=None):
        master_schema = self.schema()
        if cls:
            master_schema.context = type_dict_of(cls)
        try:
            type_of(obj) # Tests if it has a _type attribute.
        except (AttributeError, TypeError):
            return self.serialize_primitive(obj)
        else:
            return master_schema.dump(obj)

    def deserialize(self, obj, cls=None):
        master_schema = self.schema()
        # _type = serialized_type_of(obj)
        if cls:
            master_schema.context = type_dict_of(cls)
        return master_schema.load(obj)

    def serialize_primitive(self, obj):
        if isinstance(obj, Sequence) and not isinstance(obj, str):
            return [self.serialize(subobj) for subobj in obj]
        elif isinstance(obj, Mapping):
            return {key: self.serialize(value) for key, value in obj.items()}
        else:
            try:
                Field = self.base_schema_cls.TYPE_MAPPING[type(obj)]
            except KeyError:
                # No mapped field. Return as-is !
                return obj
            field = Field()
            return field._serialize(obj, None, None)

    def deserialize_primitive(self, obj, cls=None):
        if isinstance(obj, Sequence):
            return [self.deserialize(subobj, cls=cls) for subobj in obj]
        elif isinstance(obj, Mapping):
            return {key: self.deserialize(value, cls=cls) for key, value in obj.items()}
        else:
            raise ValueError('Unable to deserialize primitives other than list or dict.')


class NoopDatetime(fields.AwareDateTime):
    def _serialize(self, value, attr, obj, **kwargs):
        return value

    def _deserialize(self, value, attr, data, **kwargs):
        return value

class OrderedDictField(fields.Mapping):
    def _serialize(self, value, attr, obj, **kwargs):
        value_as_dict = dict(value)
        serialized_as_dict = super()._serialize(value_as_dict, attr, obj, **kwargs)
        return [(key, serialized_as_dict[key]) for key in value]

    def _deserialize(self, list_value, attr, data, **kwargs):
        value_as_dict = {key: value for key, value in list_value}
        deserialized_as_dict = super()._deserialize(value_as_dict, attr, data, **kwargs)
        result = OrderedDict()
        for key, _ in list_value:
            result[key] = deserialized_as_dict[key]
        return result

class FloatFieldNoNaN(fields.Float):
    def _serialize(self, value, attr, data, **kwargs):
        if value is None or isnan(value):
            value = None
        num = super()._serialize(value, attr, data, **kwargs)
        return num

class AwareDateTimeDeserializeNoop(fields.AwareDateTime):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, default_timezone=pytz.utc, **kwargs)

    def _deserialize(self, value, attr, data, **kwargs):
        if isinstance(value, datetime):
            return value.replace(tzinfo=self.default_timezone)
        else:
            return super()._deserialize(value, attr, data, **kwargs)

class DictJsonField(fields.Dict):
    def _serialize(self, value, attr, obj, **kwargs):
        return dumps(super()._serialize(value, attr, obj, **kwargs))

    def _deserialize(self, value, attr, data, **kwargs):
        return super()._deserialize(loads(value), attr, data, **kwargs)

def base_type_mapping(additions={}):
    default_type_mapping = copy(Schema.TYPE_MAPPING)
    default_type_mapping.update({
        float: FloatFieldNoNaN,
        OrderedDict: OrderedDictField,
        datetime: AwareDateTimeDeserializeNoop,
        list: fields.List,
        set: fields.List,
        tuple: fields.Tuple,
        dict: fields.Dict
    })

    try:
        from pandas import Timestamp
    except ImportError:
        pass
    else:
        default_type_mapping.update({
            Timestamp: AwareDateTimeDeserializeNoop
        })

    default_type_mapping.update(additions)
    return default_type_mapping

class AbstractSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_schemas = {}
    TYPE_MAPPING = base_type_mapping()
    type_field = TYPE_KEY
    _attached = False

    _type = fields.String(dump_only=True)

    def __new__(cls, *args, **kwargs):
        if not cls._attached:
            cls.attach_registered_dataclasses()
        obj = super().__new__(cls)
        obj.__init__(*args, **kwargs)
        return obj

    def get_obj_type(self, obj):
        # If a mapping, returns the class from the _type attribute
        # If an object, returns the class itself.
        cls = cls_of_obj(obj) 
        if cls is None:
            return None
        else:
            # Return the string registered to that cls.
            return type_of(cls)

    @classmethod
    def build_schema_type(cls, name_or_base_cls, additional_type_mapping=None, _type_mapping=None):
        if _type_mapping is None:
            if additional_type_mapping is None:
                additional_type_mapping = {}
            _type_mapping = base_type_mapping(additional_type_mapping)

        cls_dict = {
            'type_schemas': {},
            'TYPE_MAPPING': _type_mapping
        }

        if isinstance(name_or_base_cls, type):
            name = name_or_base_cls.__name__
            base_cls = (cls, name_or_base_cls)
        else:
            name = name_or_base_cls
            base_cls = (cls,)

        schema_type = type(name, base_cls, cls_dict)
        schema_type.attach_registered_dataclasses()
        return schema_type
    
    @classmethod
    def attach_registered_dataclasses(cls):
        cls.type_schemas = {}
        cls._attached = True

        for data_cls, schema_cls in _REGISTERED_TYPES:
            Schema = DataClassSchema(data_cls, schema_type=cls, schema_cls=schema_cls)

            # Building the schema now improves serialization performance by up to 800%
            cls.type_schemas[data_cls._type] = Schema()
        


def register_schema_type(key, schema_type, _default=False):
    if key in _SCHEMA_TYPES:
        raise ValueError(f'Key {key!r} is already registered.')

    _SCHEMA_TYPES[key] = schema_type

    if _default:
        _SCHEMA_TYPES[None] = schema_type

    return schema_type
    

JsonCompatSchema = AbstractSchema.build_schema_type(
    'JsonCompatSchema'
)

BsonCompatSchema = AbstractSchema.build_schema_type(
    'BsonCompatSchema',
    {
        datetime: NoopDatetime # keeps datetime as is
    }
)

DeepTabularCompatSchema = AbstractSchema.build_schema_type(
    'DeepTabularCompatSchema',
    {
        dict: DictJsonField # Serialize dicts to JSON. Required to have a constant schema.
    }
)

register_schema_type('tabular', DeepTabularCompatSchema)
register_schema_type('json', JsonCompatSchema)
register_schema_type('bson', BsonCompatSchema, _default=True)



def get_bson_serializer():
    return get_serializer('bson')

def get_json_serializer():
    return get_serializer('json')

class SerializedWrapper(dict_ops.DictWrapper):
    
    @property
    def type_(self):
        return self[TYPE_KEY]

    @type_.setter
    def type_(self, value):
        self[TYPE_KEY] = value

    def __getattr__(self, key):
        try:
            return self[key]
        except KeyError:
            raise AttributeError(key)

## JSON Schema generation

@lru_cache(maxsize=None)
def dump_json_schema(cls):
    schema = cls.get_schema()()
    json_schema_generator = JSONSchema()
    return expand_schema_refs(json_schema_generator.dump(schema))

def expand_schema_refs(json_schema, _original=None):
    if _original is None:
        _original = json_schema
    
    if '$ref' in json_schema:
        definition_path = json_schema['$ref'].split('/')[1:] # Remove first '#'
        reference_content = _original
        for p in definition_path:
            reference_content = reference_content[p]

        json_schema.update(reference_content)
        json_schema.pop('$ref')
    
    json_schema.pop('$schema', None)
    for key, value in json_schema.items():
        if isinstance(value, Mapping):
            json_schema[key] = expand_schema_refs(value, _original=_original)

    return json_schema