import os
import tempfile
from copy import copy, deepcopy

from hdmf.build import BuildManager, ObjectMapper, TypeMap
from hdmf.common.table import DynamicTable
from hdmf.container import Container, HERDManager, Data, MultiContainerInterface
from hdmf.spec import (
    AttributeSpec,
    DatasetSpec,
    DtypeSpec,
    GroupSpec,
    LinkSpec,
    NamespaceBuilder,
    NamespaceCatalog,
    RefSpec,
    SpecCatalog,
    SpecNamespace,
)
from hdmf.spec.spec import ONE_OR_MANY, ZERO_OR_MANY, ZERO_OR_ONE
from hdmf.utils import docval, get_docval, getargs

CORE_NAMESPACE = "test_core"


class CacheSpecTestHelper(object):
    @staticmethod
    def get_types(catalog):
        types = set()
        for ns_name in catalog.namespaces:
            ns = catalog.get_namespace(ns_name)
            for source in ns["schema"]:
                types.update(catalog.get_types(source["source"]))
        return types


def get_temp_filepath():
    # On Windows, h5py cannot truncate an open file in write mode.
    # The temp file will be closed before h5py truncates it and will be removed during the tearDown step.
    temp_file = tempfile.NamedTemporaryFile()
    temp_file.close()
    return temp_file.name


############################################
#  Foo example data containers and specs
###########################################
class Foo(Container):
    @docval(
        {"name": "name", "type": str, "doc": "the name of this Foo"},
        {"name": "my_data", "type": ("array_data", "data"), "doc": "some data"},
        {"name": "attr1", "type": str, "doc": "an attribute"},
        {"name": "attr2", "type": int, "doc": "another attribute"},
        {"name": "attr3", "type": float, "doc": "a third attribute", "default": 3.14},
    )
    def __init__(self, **kwargs):
        name, my_data, attr1, attr2, attr3 = getargs("name", "my_data", "attr1", "attr2", "attr3", kwargs)
        super().__init__(name=name)
        self.__data = my_data
        self.__attr1 = attr1
        self.__attr2 = attr2
        self.__attr3 = attr3

    def __eq__(self, other):
        attrs = ("name", "my_data", "attr1", "attr2", "attr3")
        return all(getattr(self, a) == getattr(other, a) for a in attrs)

    def __str__(self):
        attrs = ("name", "my_data", "attr1", "attr2", "attr3")
        return "<" + ",".join("%s=%s" % (a, getattr(self, a)) for a in attrs) + ">"

    @property
    def my_data(self):
        return self.__data

    @property
    def attr1(self):
        return self.__attr1

    @property
    def attr2(self):
        return self.__attr2

    @property
    def attr3(self):
        return self.__attr3

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


class FooBucket(Container):
    @docval(
        {"name": "name", "type": str, "doc": "the name of this bucket"},
        {"name": "foos", "type": list, "doc": "the Foo objects in this bucket", "default": list()},
    )
    def __init__(self, **kwargs):
        name, foos = getargs("name", "foos", kwargs)
        super().__init__(name=name)
        self.__foos = {f.name: f for f in foos}  # note: collections of groups are unordered in HDF5
        for f in foos:
            f.parent = self

    def __eq__(self, other):
        return self.name == other.name and self.foos == other.foos

    def __str__(self):
        return "name=%s, foos=%s" % (self.name, self.foos)

    @property
    def foos(self):
        return self.__foos

    def remove_foo(self, foo_name):
        foo = self.__foos.pop(foo_name)
        if foo.parent is self:
            self._remove_child(foo)
        return foo


class FooFile(Container, HERDManager):
    """
    NOTE: if the ROOT_NAME for the backend is not 'root' then we must set FooFile.ROOT_NAME before use
          and should be reset to 'root' when use is finished to avoid potential cross-talk between tests.
    """

    ROOT_NAME = "root"  # For HDF5 and Zarr this is the root. It should be set before use if different for the backend.

    @docval(
        {"name": "buckets", "type": list, "doc": "the FooBuckets in this file", "default": list()},
        {"name": "foo_link", "type": Foo, "doc": "an optional linked Foo", "default": None},
        {"name": "foofile_data", "type": "array_data", "doc": "an optional dataset", "default": None},
        {"name": "foo_ref_attr", "type": Foo, "doc": "a reference Foo", "default": None},
    )
    def __init__(self, **kwargs):
        buckets, foo_link = getargs("buckets", "foo_link", kwargs)
        foofile_data, foo_ref_attr = getargs("foofile_data", "foo_ref_attr", kwargs)
        super().__init__(name=self.ROOT_NAME)  # name is not used - FooFile should be the root container
        self.__buckets = {b.name: b for b in buckets}  # note: collections of groups are unordered in HDF5
        for f in buckets:
            f.parent = self
        self.__foo_link = foo_link
        self.__foofile_data = foofile_data
        self.__foo_ref_attr = foo_ref_attr

    def __eq__(self, other):
        return (
            self.buckets == other.buckets
            and self.foo_link == other.foo_link
            and self.foofile_data == other.foofile_data
        )

    def __str__(self):
        return "buckets=%s, foo_link=%s, foofile_data=%s" % (self.buckets, self.foo_link, self.foofile_data)

    @property
    def buckets(self):
        return self.__buckets

    def add_bucket(self, bucket):
        self.__buckets[bucket.name] = bucket
        bucket.parent = self

    def remove_bucket(self, bucket_name):
        bucket = self.__buckets.pop(bucket_name)
        if bucket.parent is self:
            self._remove_child(bucket)
        return bucket

    @property
    def foo_link(self):
        return self.__foo_link

    @foo_link.setter
    def foo_link(self, value):
        if self.__foo_link is None:
            self.__foo_link = value
            self.set_modified(True)
        else:
            raise ValueError("can't reset foo_link attribute")

    @property
    def foofile_data(self):
        return self.__foofile_data

    @foofile_data.setter
    def foofile_data(self, value):
        if self.__foofile_data is None:
            self.__foofile_data = value
            self.set_modified(True)
        else:
            raise ValueError("can't reset foofile_data attribute")

    @property
    def foo_ref_attr(self):
        return self.__foo_ref_attr

    @foo_ref_attr.setter
    def foo_ref_attr(self, value):
        if self.__foo_ref_attr is None:
            self.__foo_ref_attr = value
            self.set_modified(True)
        else:
            raise ValueError("can't reset foo_ref_attr attribute")


def get_foo_buildmanager(my_data_dtype="int"):
    """
    Get a BuildManager (and create all ObjectMappers) for a foofile
    :return:
    """

    foo_spec = GroupSpec(
        "A test group specification with a data type",
        data_type_def="Foo",
        datasets=[
            DatasetSpec(
                "an example dataset",
                my_data_dtype,
                name="my_data",
                shape=[None],
                attributes=[AttributeSpec("attr2", "an example integer attribute", "int")],
            )
        ],
        attributes=[
            AttributeSpec("attr1", "an example string attribute", "text"),
            AttributeSpec("attr3", "an example float attribute", "float"),
        ],
    )

    tmp_spec = GroupSpec(
        "A subgroup for Foos",
        name="foo_holder",
        groups=[
            GroupSpec(
                "the Foos in this bucket",
                data_type_inc="Foo",
                quantity=ZERO_OR_MANY,
            )
        ],
    )

    bucket_spec = GroupSpec(
        "A test group specification for a data type containing data type",
        data_type_def="FooBucket",
        groups=[tmp_spec],
    )

    class FooMapper(ObjectMapper):
        def __init__(self, spec):
            super().__init__(spec)
            my_data_spec = spec.get_dataset("my_data")
            self.map_spec("attr2", my_data_spec.get_attribute("attr2"))

    class BucketMapper(ObjectMapper):
        def __init__(self, spec):
            super().__init__(spec)
            foo_holder_spec = spec.get_group("foo_holder")
            self.unmap(foo_holder_spec)
            foo_spec = foo_holder_spec.get_data_type("Foo")
            self.map_spec("foos", foo_spec)

    file_links_spec = GroupSpec(
        "Foo link group",
        name="links",
        links=[
            LinkSpec(
                "Foo link",
                name="foo_link",
                target_type="Foo",
                quantity=ZERO_OR_ONE,
            )
        ],
    )

    file_spec = GroupSpec(
        "A file of Foos contained in FooBuckets",
        data_type_def="FooFile",
        groups=[
            GroupSpec(
                "Holds the FooBuckets",
                name="buckets",
                groups=[
                    GroupSpec(
                        "One or more FooBuckets",
                        data_type_inc="FooBucket",
                        quantity=ZERO_OR_MANY,
                    )
                ],
            ),
            file_links_spec,
        ],
        datasets=[
            DatasetSpec(
                "Foo data",
                name="foofile_data",
                dtype="int",
                quantity=ZERO_OR_ONE,
            )
        ],
        attributes=[
            AttributeSpec(
                doc="Foo ref attr",
                name="foo_ref_attr",
                dtype=RefSpec("Foo", "object"),
                required=False,
            )
        ],
    )

    class FileMapper(ObjectMapper):
        def __init__(self, spec):
            super().__init__(spec)
            bucket_spec = spec.get_group("buckets").get_data_type("FooBucket")
            self.map_spec("buckets", bucket_spec)
            self.unmap(spec.get_group("links"))
            foo_link_spec = spec.get_group("links").get_link("foo_link")
            self.map_spec("foo_link", foo_link_spec)

    spec_catalog = SpecCatalog()
    spec_catalog.register_spec(foo_spec, "test.yaml")
    spec_catalog.register_spec(bucket_spec, "test.yaml")
    spec_catalog.register_spec(file_spec, "test.yaml")
    namespace = SpecNamespace(
        "a test namespace",
        CORE_NAMESPACE,
        [{"source": "test.yaml"}],
        version="0.1.0",
        catalog=spec_catalog,
    )
    namespace_catalog = NamespaceCatalog()
    namespace_catalog.add_namespace(CORE_NAMESPACE, namespace)
    type_map = TypeMap(namespace_catalog)

    type_map.register_container_type(CORE_NAMESPACE, "Foo", Foo)
    type_map.register_container_type(CORE_NAMESPACE, "FooBucket", FooBucket)
    type_map.register_container_type(CORE_NAMESPACE, "FooFile", FooFile)

    type_map.register_map(Foo, FooMapper)
    type_map.register_map(FooBucket, BucketMapper)
    type_map.register_map(FooFile, FileMapper)

    manager = BuildManager(type_map)
    return manager


############################################
#  Baz example data containers and specs
###########################################
class Baz(Container):
    pass


class BazData(Data):
    pass


class BazCpdData(Data):
    pass


class BazBucket(Container):
    @docval(
        {"name": "name", "type": str, "doc": "the name of this bucket"},
        {"name": "bazs", "type": list, "doc": "the Baz objects in this bucket"},
        {"name": "baz_data", "type": BazData, "doc": "dataset of Baz references", "default": None},
        {"name": "baz_cpd_data", "type": BazCpdData, "doc": "dataset of Baz references", "default": None},
    )
    def __init__(self, **kwargs):
        name, bazs, baz_data, baz_cpd_data = getargs("name", "bazs", "baz_data", "baz_cpd_data", kwargs)
        super().__init__(name=name)
        self.__bazs = {b.name: b for b in bazs}  # note: collections of groups are unordered in HDF5
        for b in bazs:
            b.parent = self
        self.__baz_data = baz_data
        if self.__baz_data is not None:
            self.__baz_data.parent = self
        self.__baz_cpd_data = baz_cpd_data
        if self.__baz_cpd_data is not None:
            self.__baz_cpd_data.parent = self

    @property
    def bazs(self):
        return self.__bazs

    @property
    def baz_data(self):
        return self.__baz_data

    @property
    def baz_cpd_data(self):
        return self.__baz_cpd_data

    def add_baz(self, baz):
        self.__bazs[baz.name] = baz
        baz.parent = self

    def remove_baz(self, baz_name):
        baz = self.__bazs.pop(baz_name)
        self._remove_child(baz)
        return baz


def get_baz_buildmanager():
    baz_spec = GroupSpec(
        doc="A test group specification with a data type",
        data_type_def="Baz",
    )

    baz_data_spec = DatasetSpec(
        doc="A test dataset of references specification with a data type",
        name="baz_data",
        data_type_def="BazData",
        dtype=RefSpec("Baz", "object"),
        shape=[None],
    )

    baz_cpd_data_spec = DatasetSpec(
        doc="A test compound dataset with references specification with a data type",
        name="baz_cpd_data",
        data_type_def="BazCpdData",
        dtype=[
            DtypeSpec(name="part1", doc="doc", dtype="int"),
            DtypeSpec(name="part2", doc="doc", dtype=RefSpec("Baz", "object")),
        ],
        shape=[None],
    )

    baz_holder_spec = GroupSpec(
        doc="group of bazs",
        name="bazs",
        groups=[GroupSpec(doc="Baz", data_type_inc="Baz", quantity=ONE_OR_MANY)],
    )

    baz_bucket_spec = GroupSpec(
        doc="A test group specification for a data type containing data type",
        data_type_def="BazBucket",
        groups=[baz_holder_spec],
        datasets=[
            DatasetSpec(doc="doc", data_type_inc="BazData", quantity=ZERO_OR_ONE),
            DatasetSpec(doc="doc", data_type_inc="BazCpdData", quantity=ZERO_OR_ONE),
        ],
    )

    spec_catalog = SpecCatalog()
    spec_catalog.register_spec(baz_spec, "test.yaml")
    spec_catalog.register_spec(baz_data_spec, "test.yaml")
    spec_catalog.register_spec(baz_cpd_data_spec, "test.yaml")
    spec_catalog.register_spec(baz_bucket_spec, "test.yaml")

    namespace = SpecNamespace(
        "a test namespace",
        CORE_NAMESPACE,
        [{"source": "test.yaml"}],
        version="0.1.0",
        catalog=spec_catalog,
    )

    namespace_catalog = NamespaceCatalog()
    namespace_catalog.add_namespace(CORE_NAMESPACE, namespace)

    type_map = TypeMap(namespace_catalog)
    type_map.register_container_type(CORE_NAMESPACE, "Baz", Baz)
    type_map.register_container_type(CORE_NAMESPACE, "BazData", BazData)
    type_map.register_container_type(CORE_NAMESPACE, "BazCpdData", BazCpdData)
    type_map.register_container_type(CORE_NAMESPACE, "BazBucket", BazBucket)

    class BazBucketMapper(ObjectMapper):
        def __init__(self, spec):
            super().__init__(spec)
            baz_holder_spec = spec.get_group("bazs")
            self.unmap(baz_holder_spec)
            baz_spec = baz_holder_spec.get_data_type("Baz")
            self.map_spec("bazs", baz_spec)

    type_map.register_map(BazBucket, BazBucketMapper)

    manager = BuildManager(type_map)
    return manager


def create_test_type_map(specs, container_classes, mappers=None):
    """
    Create a TypeMap with the specs registered under a test namespace, and classes and mappers registered to type names.
    :param specs: list of specs
    :param container_classes: dict of type name to container class
    :param mappers: (optional) dict of type name to mapper class
    :return: the constructed TypeMap
    """
    spec_catalog = SpecCatalog()
    schema_file = "test.yaml"
    for s in specs:
        spec_catalog.register_spec(s, schema_file)
    namespace = SpecNamespace(
        doc="a test namespace",
        name=CORE_NAMESPACE,
        schema=[{"source": schema_file}],
        version="0.1.0",
        catalog=spec_catalog,
    )
    namespace_catalog = NamespaceCatalog()
    namespace_catalog.add_namespace(CORE_NAMESPACE, namespace)
    type_map = TypeMap(namespace_catalog)
    for type_name, container_cls in container_classes.items():
        type_map.register_container_type(CORE_NAMESPACE, type_name, container_cls)
    if mappers:
        for type_name, mapper_cls in mappers.items():
            container_cls = container_classes[type_name]
            type_map.register_map(container_cls, mapper_cls)
    return type_map


def create_load_namespace_yaml(namespace_name, specs, output_dir, incl_types, type_map):
    """
    Create a TypeMap with the specs loaded from YAML files and dependencies resolved.

    This writes namespaces and specs to YAML files, creates an empty TypeMap, and calls load_namespaces on the
    TypeMap, instead of manually creating a SpecCatalog, SpecNamespace, NamespaceCatalog and manually registering
    container types. Importantly, this process resolves dependencies across namespaces.

    :param namespace_name: Name of the new namespace.
    :param specs: List of specs of new data types to add.
    :param incl_types: Dict mapping included namespace name to list of data types to include or None to include all.
    :param type_map: The type map to load the namespace into.
    """
    ns_builder = NamespaceBuilder(
        name=namespace_name,
        doc="a test namespace",
        version="0.1.0",
    )
    ns_filename = ns_builder.name + ".namespace.yaml"
    ext_filename = ns_builder.name + ".extensions.yaml"

    for ns, types in incl_types.items():
        if types is None:  # include all types
            ns_builder.include_namespace(ns)
        else:
            for dt in types:
                ns_builder.include_type(dt, namespace=ns)

    for data_type in specs:
        ns_builder.add_spec(ext_filename, data_type)

    ns_builder.export(ns_filename, outdir=output_dir)
    ns_path = os.path.join(output_dir, ns_filename)
    type_map.load_namespaces(ns_path)


# ##### custom spec classes #####


def swap_inc_def(cls, custom_cls):
    args = get_docval(cls.__init__)
    ret = list()
    for arg in args:
        if arg["name"] == "data_type_def":
            ret.append(
                {
                    "name": "my_data_type_def",
                    "type": str,
                    "doc": "the NWB data type this spec defines",
                    "default": None,
                }
            )
        elif arg["name"] == "data_type_inc":
            ret.append(
                {
                    "name": "my_data_type_inc",
                    "type": (custom_cls, str),
                    "doc": "the NWB data type this spec includes",
                    "default": None,
                }
            )
        else:
            ret.append(copy(arg))
    return ret


class BaseStorageOverride:
    __type_key = "my_data_type"
    __inc_key = "my_data_type_inc"
    __def_key = "my_data_type_def"

    @classmethod
    def type_key(cls):
        """Get the key used to store data type on an instance"""
        return cls.__type_key

    @classmethod
    def inc_key(cls):
        """Get the key used to define a data_type include."""
        return cls.__inc_key

    @classmethod
    def def_key(cls):
        """Get the key used to define a data_type definition."""
        return cls.__def_key

    @classmethod
    def build_const_args(cls, spec_dict):
        """Extend base functionality to remap data_type_def and data_type_inc keys"""
        spec_dict = copy(spec_dict)
        proxy = super(BaseStorageOverride, cls)
        if proxy.inc_key() in spec_dict:
            spec_dict[cls.inc_key()] = spec_dict.pop(proxy.inc_key())
        if proxy.def_key() in spec_dict:
            spec_dict[cls.def_key()] = spec_dict.pop(proxy.def_key())
        ret = proxy.build_const_args(spec_dict)
        return ret

    @classmethod
    def _translate_kwargs(cls, kwargs):
        """Swap mydata_type_def and mydata_type_inc for data_type_def and data_type_inc, respectively"""
        proxy = super(BaseStorageOverride, cls)
        kwargs[proxy.def_key()] = kwargs.pop(cls.def_key())
        kwargs[proxy.inc_key()] = kwargs.pop(cls.inc_key())
        return kwargs


class CustomGroupSpec(BaseStorageOverride, GroupSpec):
    @docval(*deepcopy(swap_inc_def(GroupSpec, "CustomGroupSpec")))
    def __init__(self, **kwargs):
        kwargs = self._translate_kwargs(kwargs)
        super().__init__(**kwargs)

    @classmethod
    def dataset_spec_cls(cls):
        return CustomDatasetSpec

    @docval(*deepcopy(swap_inc_def(GroupSpec, "CustomGroupSpec")))
    def add_group(self, **kwargs):
        spec = CustomGroupSpec(**kwargs)
        self.set_group(spec)
        return spec

    @docval(*deepcopy(swap_inc_def(DatasetSpec, "CustomDatasetSpec")))
    def add_dataset(self, **kwargs):
        """Add a new specification for a subgroup to this group specification"""
        spec = CustomDatasetSpec(**kwargs)
        self.set_dataset(spec)
        return spec


class CustomDatasetSpec(BaseStorageOverride, DatasetSpec):
    @docval(*deepcopy(swap_inc_def(DatasetSpec, "CustomDatasetSpec")))
    def __init__(self, **kwargs):
        kwargs = self._translate_kwargs(kwargs)
        super().__init__(**kwargs)


class CustomSpecNamespace(SpecNamespace):
    __types_key = "my_data_types"

    @classmethod
    def types_key(cls):
        return cls.__types_key

class FooExtendDynamicTable0(DynamicTable):
    """
    Within PyNWB, PlaneSegmentation extends DynamicTable and sets __columns__. This class is a helper
    class for testing and is directly meant to test __gather_columns, i.e., class generation, downstream.
    """
    __columns__ = (
        {'name': 'col1', 'description': '...'},
        {'name': 'col2', 'description': '...'},
    )

    def __init__(self, **kwargs):
        kwargs['name'] = 'foo0'
        kwargs['description'] = '...'
        super().__init__(**kwargs)


class FooExtendDynamicTable1(FooExtendDynamicTable0):
    """
    In extensions, users can create new classes that inherit from classes that inherit from DynamicTable.
    This is a helper class for testing and is directly meant to test __gather_columns, i.e.,
    class generation, downstream.
    """
    __columns__ = (
        {'name': 'col3', 'description': '...'},
        {'name': 'col4', 'description': '...'},
    )

    def __init__(self, **kwargs):
        kwargs['name'] = 'foo1'
        kwargs['description'] = '...'
        super().__init__(**kwargs)


class FooExtendDynamicTable2(FooExtendDynamicTable1, MultiContainerInterface):
    __clsconf__ = {
        'add': '...',
        'get': '...',
        'create': '...',
        'type': Container,
        'attr': '...'
    }

    def __init__(self, **kwargs):
        kwargs['name'] = 'foo2'
        kwargs['description'] = '...'
        super().__init__(**kwargs)