from six.moves.cPickle import dumps, load, loads, PickleError
import warnings
import io
import sys

from pyface.qt import QtCore


#-------------------------------------------------------------------------------
#  'PyMimeData' class:
#-------------------------------------------------------------------------------

if sys.version_info[0] < 3:
    def str2bytes(s):
        return s
else:
    def str2bytes(s):
        return bytes(s,'ascii')


class PyMimeData(QtCore.QMimeData):
    """ The PyMimeData wraps a Python instance as MIME data.
    """
    # The MIME type for instances.
    MIME_TYPE = u'application/x-ets-qt4-instance'
    NOPICKLE_MIME_TYPE = u'application/x-ets-qt4-instance-no-pickle'

    def __init__(self, data=None, pickle=True):
        """ Initialise the instance.
        """
        QtCore.QMimeData.__init__(self)

        # Keep a local reference to be returned if possible.
        self._local_instance = data

        if pickle:
            if data is not None:
                # We may not be able to pickle the data.
                try:
                    pdata = dumps(data)
                    # This format (as opposed to using a single sequence) allows
                    # the type to be extracted without unpickling the data.
                    self.setData(self.MIME_TYPE, dumps(data.__class__) + pdata)
                except (PickleError, TypeError, AttributeError):
                    # if pickle fails, still try to create a draggable
                    warnings.warn(("Could not pickle dragged object %s, " +
                            "using %s mimetype instead") % (repr(data),
                            self.NOPICKLE_MIME_TYPE), RuntimeWarning)
                    self.setData(self.NOPICKLE_MIME_TYPE, str2bytes(str(id(data))))

        else:
            self.setData(self.NOPICKLE_MIME_TYPE, str2bytes(str(id(data))))

    @classmethod
    def coerce(cls, md):
        """ Wrap a QMimeData or a python object to a PyMimeData.
        """
        # See if the data is already of the right type.  If it is then we know
        # we are in the same process.
        if isinstance(md, cls):
            return md

        if isinstance(md, PyMimeData):
            # if it is a PyMimeData, migrate all its data, subclasses should
            # override this method if it doesn't do thgs correctly for them
            data = md.instance()
            nmd = cls()
            nmd._local_instance = data
            for format in md.formats():
                nmd.setData(format, md.data(format))
        elif isinstance(md, QtCore.QMimeData):
            # if it is a QMimeData, migrate all its data
            nmd = cls()
            for format in md.formats():
                nmd.setData(format, md.data(format))
        else:
            # by default, try to pickle the coerced object
            pickle = True

            # See if the data is a list, if so check for any items which are
            # themselves of the right type.  If so, extract the instance and
            # track whether we should pickle.
            # XXX lists should suffice for now, but may want other containers
            if isinstance(md, list):
                pickle = not any(item.hasFormat(cls.NOPICKLE_MIME_TYPE)
                        for item in md if isinstance(item, QtCore.QMimeData))
                md = [item.instance() if isinstance(item, PyMimeData) else item
                        for item in md]

            # Arbitrary python object, wrap it into PyMimeData
            nmd = cls(md, pickle)

        return nmd

    def instance(self):
        """ Return the instance.
        """
        if self._local_instance is not None:
            return self._local_instance

        if not self.hasFormat(self.MIME_TYPE):
            # We have no pickled python data defined.
            return None

        stream = io.BytesIO(self.data(self.MIME_TYPE).data())

        try:
            # Skip the type.
            load(stream)

            # Recreate the instance.
            return load(stream)
        except PickleError:
            pass

        return None

    def instanceType(self):
        """ Return the type of the instance.
        """
        if self._local_instance is not None:
            return self._local_instance.__class__

        try:
            if self.hasFormat(self.MIME_TYPE):
                return loads(self.data(self.MIME_TYPE).data())
        except PickleError:
            pass

        return None

    def localPaths(self):
        """ The list of local paths from url list, if any.
        """
        ret = []
        for url in self.urls():
            if url.scheme() == 'file':
                ret.append(url.toLocalFile())
        return ret