#
# This file is distributed under the MIT License. See LICENSE.md for details.
#

import copy
from abc import ABC
from typing import Dict, List

from .definition import Definition
from .enum import EnumDefinition
from .reference import ReferenceDefinition
from .sequence import SequenceDefinition
from .upcastable import UpcastableDefinition


def is_enum(resolved_type):
    return isinstance(resolved_type, EnumDefinition)


class StructField(ABC):
    def __init__(self, *, name, doc=None, const=False, is_guid=False):
        self.name = name
        self.doc = doc
        self.is_key = False
        self.upcastable = False
        self.const = const
        self.is_guid = is_guid
        self.resolved_type = None

    @staticmethod
    def from_yaml(source_dict: Dict):
        if source_dict.get("type"):
            return SimpleStructField(**source_dict)

        if source_dict.get("sequence"):
            sequence_def = source_dict["sequence"]
            args = copy.copy(source_dict)
            del args["sequence"]
            args.update(
                {
                    "sequence_type": sequence_def["type"],
                    "element_type": sequence_def["elementType"],
                }
            )
            return SequenceStructField(**args)

        if source_dict.get("referenceTo"):
            args = copy.copy(source_dict)
            args["pointee_type"] = source_dict["referenceTo"]
            del args["referenceTo"]
            return ReferenceStructField(**args)

        raise ValueError("Invalid struct field")

    def resolve_references(self, schema):
        raise NotImplementedError()


class SimpleStructField(StructField):
    def __init__(
        self,
        *,
        name,
        type,  # noqa: A002
        doc=None,
        const=False,
        default=None,
        is_guid=False,
    ):
        super().__init__(name=name, doc=doc, const=const, is_guid=is_guid)
        self.type = type
        self.default = default

    def resolve_references(self, schema):
        self.resolved_type = schema.get_definition_for(self.type)

        if isinstance(self.resolved_type, StructDefinition) and self.resolved_type.abstract:
            self.resolved_type = UpcastableDefinition(self.resolved_type)
            self.upcastable = True

        assert self.resolved_type


class SequenceStructField(StructField):
    def __init__(
        self,
        *,
        name,
        sequence_type,
        element_type,
        doc=None,
        const=False,
    ):
        super().__init__(name=name, doc=doc, const=const)
        self.sequence_type = sequence_type
        self.element_type = element_type
        self.resolved_element_type = None

    def resolve_references(self, schema):
        self.resolved_element_type = schema.get_definition_for(self.element_type)
        if (
            isinstance(self.resolved_element_type, StructDefinition)
            and self.resolved_element_type.abstract
        ):
            self.resolved_element_type = UpcastableDefinition(self.resolved_element_type)
            self.upcastable = True
        self.resolved_type = SequenceDefinition(self.sequence_type, self.resolved_element_type)


class ReferenceStructField(StructField):
    def __init__(self, *, name, pointee_type, doc=None, const=False):
        super().__init__(name=name, doc=doc, const=const)
        self.pointee_type = pointee_type

    def resolve_references(self, schema):
        resolved_pointee_type = schema.get_definition_for(self.pointee_type)
        resolved_root_type = schema.get_definition_for(schema.root_type)
        self.resolved_type = ReferenceDefinition(resolved_pointee_type, resolved_root_type)


class StructDefinition(Definition):
    def __init__(
        self,
        *,
        name,
        fields,
        inherits=None,
        doc=None,
        abstract=False,
        _key=None,
    ):
        super().__init__(False, name)

        self.doc = doc
        # Names of types on which this definition depends on.
        # Not necessarily names defined by the user
        self.autogenerated = False

        self.fields: List[StructField] = fields
        self.abstract = abstract
        self._key = _key
        self._inherits = inherits

        # These fields will be populated by resolve_references()
        self.inherits = None
        self._key_kind_index = None
        # None, "simple" or "composite"
        self.keytype = None
        self.key_fields: List[SimpleStructField] = []
        self.emit_full_constructor = None
        self.children: List[StructDefinition] = []
        self._refs_resolved = False

    def resolve_references(self, schema):
        if self._refs_resolved:
            return

        if self._inherits:
            self.dependencies.add(self._inherits)

            self.inherits = schema.get_definition_for(self._inherits)
            self.inherits.resolve_references(schema)

            # TODO: support indirect inheritance
            assert not self._key
            self._key = self.inherits._key
            self._key_kind_index = self.inherits._key_kind_index
            self.keytype = self.inherits.keytype
            self.key_fields = self.inherits.key_fields

        for field in self.fields:
            field.resolve_references(schema)

            if isinstance(field, SimpleStructField):
                self.dependencies.add(field.type)
            elif isinstance(field, SequenceStructField):
                self.dependencies.add(field.element_type)
            elif isinstance(field, ReferenceStructField):
                self.dependencies.add(field.pointee_type)
                # TODO: technically, every reference also depends on the root
                # type, but adding it as a dependency here will cause a circular
                # dependency.
                #
                # self.dependencies.add(schema.root_type)  # noqa: E800
            else:
                raise ValueError()

        if self._key and not self._inherits:
            for field in self.fields:
                if field.name in self._key:
                    assert isinstance(field, SimpleStructField)

                    field.is_key = True
                    # TODO: mark `field` as `const`

                    self.key_fields.append(field)

            kind_fields = [i for i, k in enumerate(self._key) if k == "Kind"]
            assert len(kind_fields) <= 1, f"Multiple kind fields in {self.name}"
            self._key_kind_index = kind_fields[0] if kind_fields else None

            if len(self.key_fields) == 0:
                self.keytype = None
            elif len(self.key_fields) == 1:
                self.keytype = "simple"
            else:
                self.keytype = "composite"

        self._refs_resolved = True

        self.emit_full_constructor = len(set(self.all_optional_fields)) != 0

        if (self.inherits or self.abstract) and self.keytype:
            assert self._key_kind_index is not None, "A polymorphic type without kind in the key"

    @staticmethod
    def from_dict(source_dict: Dict):
        args = {
            "name": source_dict["name"],
            "doc": source_dict.get("doc"),
            "fields": [StructField.from_yaml(d) for d in source_dict["fields"]],
            "_key": source_dict.get("key"),
            "inherits": source_dict.get("inherits"),
            "abstract": source_dict.get("abstract", False),
        }
        return StructDefinition(**args)

    def is_required(self, field):
        if field.is_key:
            return True

        return (self.inherits or self.abstract) and field.name == "Kind"

    @property
    def optional_fields(self):
        for field in self.fields:
            if not self.is_required(field):
                yield field

    @property
    def required_fields(self):
        for field in self.fields:
            if self.is_required(field):
                yield field

    @property
    def all_optional_fields(self):
        for field in self.all_fields:
            if not self.is_required(field):
                yield field

    @property
    def all_required_fields(self):
        for field in self.all_fields:
            if self.is_required(field):
                yield field

    @property
    def all_fields(self):
        if not self._refs_resolved:
            raise RuntimeError("`all_fields` is only available after resolving references")
        if self.inherits:
            yield from self.inherits.fields
        yield from self.fields