from datetime import datetime
from dateutil.tz import tzlocal, tzutc
import os
from abc import ABCMeta, abstractmethod
import warnings

from pynwb import NWBFile, NWBHDF5IO, validate as pynwb_validate
from .utils import remove_test_file
from hdmf.backends.warnings import BrokenLinkWarning
from hdmf.build.warnings import MissingRequiredWarning, OrphanContainerWarning


class NWBH5IOMixin(metaclass=ABCMeta):
    """
    Mixin class for methods to run a roundtrip test writing an NWB file with an Container and reading the Container
    from the NWB file. The setUp, test_roundtrip, and tearDown methods will be run by unittest.

    The abstract methods setUpContainer, addContainer, and getContainer needs to be implemented by classes that include
    this mixin.

    Example::

        class TestMyContainerIO(NWBH5IOMixin, TestCase):
            def setUpContainer(self):
                # return a test Container to read/write
            def addContainer(self, nwbfile):
                # add the test Container to an NWB file
            def getContainer(self, nwbfile):
                # return the test Container from an NWB file

    This code is adapted from hdmf.testing.H5RoundTripMixin.
    """

    def setUp(self):
        self.container = self.setUpContainer()
        self.start_time = datetime(1971, 1, 1, 12, tzinfo=tzutc())
        self.create_date = datetime(2018, 4, 15, 12, tzinfo=tzlocal())
        self.container_type = self.container.__class__.__name__
        self.filename = 'test_%s.nwb' % self.container_type
        self.export_filename = 'test_export_%s.nwb' % self.container_type
        self.writer = None
        self.reader = None
        self.export_reader = None

    def tearDown(self):
        if self.writer is not None:
            self.writer.close()
        if self.reader is not None:
            self.reader.close()
        if self.export_reader is not None:
            self.export_reader.close()
        remove_test_file(self.filename)
        remove_test_file(self.export_filename)

    @abstractmethod
    def setUpContainer(self):
        """ Should return the test Container to read/write """
        raise NotImplementedError('Cannot run test unless setUpContainer is implemented')

    def test_roundtrip(self):
        """
        Test whether the test Container read from file has the same contents as the original test Container and
        validate the file
        """
        self.read_container = self.roundtripContainer()
        self.assertIsNotNone(str(self.container))  # added as a test to make sure printing works
        self.assertIsNotNone(str(self.read_container))
        # make sure we get a completely new object
        self.assertNotEqual(id(self.container), id(self.read_container))
        # make sure the object ID is preserved
        self.assertIs(self.read_nwbfile.objects[self.container.object_id], self.read_container)
        self.assertContainerEqual(self.read_container, self.container)

    def test_roundtrip_export(self):
        """
        Test whether the test Container read from an exported file has the same contents as the original test Container
        and validate the file
        """
        self.read_container = self.roundtripExportContainer()
        self.assertIsNotNone(str(self.read_container))  # added as a test to make sure printing works
        # make sure we get a completely new object
        self.assertNotEqual(id(self.container), id(self.read_container))
        # make sure the object ID is preserved
        self.assertIs(self.read_exported_nwbfile.objects[self.container.object_id], self.read_container)
        self.assertContainerEqual(self.read_container, self.container, ignore_hdmf_attrs=True)

    def roundtripContainer(self, cache_spec=False):
        """
        Add the test Container to an NWBFile, write it to file, read the file, and return the test Container from the
        file
        """
        description = 'a file to test writing and reading a %s' % self.container_type
        identifier = 'TEST_%s' % self.container_type
        nwbfile = NWBFile(description, identifier, self.start_time, file_create_date=self.create_date)
        self.addContainer(nwbfile)

        with warnings.catch_warnings(record=True) as ws:
            with NWBHDF5IO(self.filename, mode='w') as write_io:
                write_io.write(nwbfile, cache_spec=cache_spec)

            self.validate()

            self.reader = NWBHDF5IO(self.filename, mode='r')
            self.read_nwbfile = self.reader.read()

        if ws:
            for w in ws:
                if issubclass(w.category, (MissingRequiredWarning,
                                           OrphanContainerWarning,
                                           BrokenLinkWarning)):
                    raise Exception('%s: %s' % (w.category.__name__, w.message))
                else:
                    warnings.warn(w.message, w.category)

        try:
            return self.getContainer(self.read_nwbfile)
        except Exception as e:
            self.reader.close()
            self.reader = None
            raise e

    def roundtripExportContainer(self, cache_spec=False):
        """
        Add the test Container to an NWBFile, write it to file, read the file, export the read NWBFile to another
        file, and return the test Container from the file
        """
        self.roundtripContainer(cache_spec=cache_spec)  # self.read_nwbfile is now set

        with warnings.catch_warnings(record=True) as ws:
            NWBHDF5IO.export_io(
                src_io=self.reader,
                path=self.export_filename,
                cache_spec=cache_spec,
            )

            self.validate()

            self.export_reader = NWBHDF5IO(self.export_filename, mode='r')
            self.read_exported_nwbfile = self.export_reader.read()

        if ws:
            for w in ws:
                if issubclass(w.category, (MissingRequiredWarning,
                                           OrphanContainerWarning,
                                           BrokenLinkWarning)):
                    raise Exception('%s: %s' % (w.category.__name__, w.message))
                else:
                    warnings.warn(w.message, w.category)

        try:
            return self.getContainer(self.read_exported_nwbfile)
        except Exception as e:
            self.export_reader.close()
            self.export_reader = None
            raise e

    @abstractmethod
    def addContainer(self, nwbfile):
        """ Should add the test Container to the given NWBFile """
        raise NotImplementedError('Cannot run test unless addContainer is implemented')

    @abstractmethod
    def getContainer(self, nwbfile):
        """ Should return the test Container from the given NWBFile """
        raise NotImplementedError('Cannot run test unless getContainer is implemented')

    def validate(self):
        """ Validate the created files """
        if os.path.exists(self.filename):
            with NWBHDF5IO(self.filename, mode='r') as io:
                errors = pynwb_validate(io)
                if errors:
                    for err in errors:
                        raise Exception(err)

        if os.path.exists(self.export_filename):
            with NWBHDF5IO(self.filename, mode='r') as io:
                errors = pynwb_validate(io)
                if errors:
                    for err in errors:
                        raise Exception(err)


class AcquisitionH5IOMixin(NWBH5IOMixin):
    """
    Mixin class for methods to run a roundtrip test writing an NWB file with an Container as an acquisition and reading
    the Container as an acquisition from the NWB file. The setUp, test_roundtrip, and tearDown methods will be run by
    unittest.

    The abstract method setUpContainer needs to be implemented by classes that include this mixin.

    Example::

        class TestMyContainerIO(NWBH5IOMixin, TestCase):
            def setUpContainer(self):
                # return a test Container to read/write

    This code is adapted from hdmf.testing.H5RoundTripMixin.
    """

    def addContainer(self, nwbfile):
        ''' Add an NWBDataInterface object to the file as an acquisition '''
        nwbfile.add_acquisition(self.container)

    def getContainer(self, nwbfile):
        ''' Get the NWBDataInterface object from the file '''
        return nwbfile.get_acquisition(self.container.name)
