import difflib
import os
import io
import shutil
import struct
import sys
import tarfile
import tempfile
import warnings
from contextlib import closing, contextmanager
# from ._utils import _import_dotted_name
# from ._six import string_classes as _string_classes
from typing import Any, BinaryIO, cast, Dict, Optional, Type, Tuple, Union, IO
import copyreg
import pickle
import pathlib
from py7zr import SevenZipFile, is_7zfile

# Protocol version 0 is the original “human-readable” protocol and is backwards compatible with earlier versions of Python.
# Protocol version 1 is an old binary format which is also compatible with earlier versions of Python.
# Protocol version 2 was introduced in Python 2.3. It provides much more efficient pickling of new-style classes. Refer to PEP 307 for information about improvements brought by protocol 2.
# Protocol version 3 was added in Python 3.0. It has explicit support for bytes objects and cannot be unpickled by Python 2.x. This was the default protocol in Python 3.0–3.7.
# Protocol version 4 was added in Python 3.4. It adds support for very large objects, pickling more kinds of objects, and some data format optimizations. It is the default protocol starting with Python 3.8. Refer to PEP 3154 for information about improvements brought by protocol 4.
# Protocol version 5 was added in Python 3.8. It adds support for out-of-band data and speedup for in-band data. Refer to PEP 574 for information about improvements brought by protocol 5.
import sys
if sys.version.startswith('2'):
    DEFAULT_PROTOCOL = 2
else:
    DEFAULT_PROTOCOL = 3

LONG_SIZE = struct.Struct('=l').size
INT_SIZE = struct.Struct('=i').size
SHORT_SIZE = struct.Struct('=h').size


MAGIC_NUMBER = 0x5448554e45543130
PROTOCOL_VERSION = 1001
STORAGE_KEY_SEPARATOR = ','

class SourceChangeWarning(Warning):
    pass


@contextmanager
def mkdtemp():
    path = tempfile.mkdtemp()
    yield path
    shutil.rmtree(path)


_package_registry = []


def _is_zipfile(f) -> bool:
    # This is a stricter implementation than zipfile.is_zipfile().
    # zipfile.is_zipfile() is True if the magic number appears anywhere in the
    # binary. Since we expect the files here to be generated by torch.save or
    # torch.jit.save, it's safe to only check the start bytes and avoid
    # collisions and assume the zip has only 1 file.
    # See bugs.python.org/issue28494.

    # Read the first 4 bytes of the file
    read_bytes = []
    start = f.tell()

    byte = f.read(1)
    while byte != "":
        read_bytes.append(byte)
        if len(read_bytes) == 4:
            break
        byte = f.read(1)
    f.seek(start)

    local_header_magic_number = [b'P', b'K', b'\x03', b'\x04']
    return read_bytes == local_header_magic_number


def register_package(priority, tagger, deserializer):
    queue_elem = (priority, tagger, deserializer)
    _package_registry.append(queue_elem)
    _package_registry.sort()


def check_module_version_greater_or_equal(module, req_version_tuple, error_if_malformed=True):
    '''
    Check if a module's version satisfies requirements

    Usually, a module's version string will be like 'x.y.z', which would be represented
    as a tuple (x, y, z), but sometimes it could be an unexpected format. If the version
    string does not match the given tuple's format up to the length of the tuple, then
    error and exit or emit a warning.

    Args:
        module: the module to check the version of
        req_version_tuple: tuple (usually of ints) representing the required version
        error_if_malformed: whether we should exit if module version string is malformed

    Returns:
        requirement_is_met: bool
    '''
    try:
        version_strs = module.__version__.split('.')
        # Cast module version fields to match the types of the required version
        module_version = tuple(
            type(req_field)(version_strs[idx]) for idx, req_field in enumerate(req_version_tuple)
        )
        requirement_is_met = module_version >= req_version_tuple

    except Exception as e:
        message = (
            "'%s' module version string is malformed '%s' and cannot be compared"
            " with tuple %s"
        ) % (
            module.__name__, module.__version__, str(req_version_tuple)
        )
        if error_if_malformed:
            raise RuntimeError(message) from e
        else:
            warnings.warn(message + ', but continuing assuming that requirement is met')
            requirement_is_met = True

    return requirement_is_met


def _cpu_tag(obj):
    if obj.device.type == 'cpu':
        return 'cpu'


def _cuda_tag(obj):
    if obj.device.type == 'cuda':
        return 'cuda:' + str(obj.device.index)


def _cpu_deserialize(obj, location):
    if location == 'cpu':
        return obj


register_package(10, _cpu_tag, _cpu_deserialize)


def _is_path(name_or_buffer):
    return isinstance(name_or_buffer, str) or \
        isinstance(name_or_buffer, pathlib.Path)


class _opener(object):
    def __init__(self, file_like):
        self.file_like = file_like

    def __enter__(self):
        return self.file_like

    def __exit__(self, *args):
        pass


class _open_file(_opener):
    def __init__(self, name, mode):
        super(_open_file, self).__init__(open(name, mode))

    def __exit__(self, *args):
        self.file_like.close()


class _open_buffer_reader(_opener):
    def __init__(self, buffer):
        super(_open_buffer_reader, self).__init__(buffer)
        _check_seekable(buffer)


class _open_buffer_writer(_opener):
    def __exit__(self, *args):
        self.file_like.flush()


def _open_file_like(name_or_buffer, mode):
    if _is_path(name_or_buffer):
        return _open_file(name_or_buffer, mode)
    else:
        if 'w' in mode:
            return _open_buffer_writer(name_or_buffer)
        elif 'r' in mode:
            return _open_buffer_reader(name_or_buffer)
        else:
            raise RuntimeError(f"Expected 'r' or 'w' in mode but got {mode}")


class _open_zipfile_reader(_opener):
    def __init__(self, name_or_buffer) -> None:
        super(_open_zipfile_reader, self).__init__(SevenZipFile(name_or_buffer))


class _open_zipfile_writer_file(_opener):
    def __init__(self, name) -> None:
        super(_open_zipfile_writer_file, self).__init__(SevenZipFile(str(name), mode='w'))

    def __exit__(self, *args) -> None:
        self.file_like.close()


class _open_zipfile_writer_buffer(_opener):
    def __init__(self, buffer) -> None:
        self.buffer = buffer
        super(_open_zipfile_writer_buffer, self).__init__(SevenZipFile(buffer))

    def __exit__(self, *args) -> None:
        self.file_like.close()
        self.buffer.flush()


def _open_zipfile_writer(name_or_buffer):
    container: Type[_opener]
    if _is_path(name_or_buffer):
        container = _open_zipfile_writer_file
    else:
        container = _open_zipfile_writer_buffer
    return container(name_or_buffer)


def _is_compressed_file(f) -> bool:
    compress_modules = ['gzip']
    try:
        return f.__module__ in compress_modules
    except AttributeError:
        return False


def _should_read_directly(f):
    """
    Checks if f is a file that should be read directly. It should be read
    directly if it is backed by a real file (has a fileno) and is not a
    a compressed file (e.g. gzip)
    """
    if _is_compressed_file(f):
        return False
    try:
        return f.fileno() >= 0
    except io.UnsupportedOperation:
        return False
    except AttributeError:
        return False


def _check_seekable(f) -> bool:

    def raise_err_msg(patterns, e):
        for p in patterns:
            if p in str(e):
                msg = (str(e) + ". You can only torch.load from a file that is seekable."
                                + " Please pre-load the data into a buffer like io.BytesIO and"
                                + " try to load from it instead.")
                raise type(e)(msg)
        raise e

    try:
        f.seek(f.tell())
        return True
    except (io.UnsupportedOperation, AttributeError) as e:
        raise_err_msg(["seek", "tell"], e)
    return False

def _check_dill_version(pickle_module) -> None:
    '''Checks if using dill as the pickle module, and if so, checks if it is the correct version.
    If dill version is lower than 0.3.1, a ValueError is raised.

    Args:
        pickle_module: module used for pickling metadata and objects

    '''
    if pickle_module.__name__ == 'dill':
        required_dill_version = (0, 3, 1)
        if not check_module_version_greater_or_equal(pickle_module, required_dill_version, False):
            raise ValueError((
                "'torch' supports dill >= %s, but you have dill %s."
                " Please upgrade dill or switch to 'pickle'"
            ) % (
                '.'.join([str(num) for num in required_dill_version]),
                pickle_module.__version__
            ))

def save(obj, f: Union[str, os.PathLike, BinaryIO, IO[bytes]],
         pickle_module=pickle, pickle_protocol=DEFAULT_PROTOCOL, _use_new_zipfile_serialization=True) -> None:
    if _use_new_zipfile_serialization:
        with _open_zipfile_writer(f) as opened_zipfile:
            data_buf = io.BytesIO()
            pickler = pickle_module.Pickler(data_buf, protocol=pickle_protocol)
            # pickler.persistent_id = persistent_id
            pickler.dump(obj)
            read_data_buf = io.BytesIO(data_buf.getvalue())
            opened_zipfile.writef(read_data_buf, "data.pkl")
            del data_buf, read_data_buf
    else:
        with _open_file_like(f, 'wb') as opened_file:
            f = opened_file
            pickle_module.dump(MAGIC_NUMBER, f, protocol=pickle_protocol)
            pickle_module.dump(PROTOCOL_VERSION, f, protocol=pickle_protocol)
            pickle_module.dump({}, f, protocol=pickle_protocol)
            pickler = pickle_module.Pickler(f, protocol=pickle_protocol)
            # pickler.persistent_id = persistent_id
            pickler.dump(obj)

            # serialized_storage_keys = sorted(serialized_storages.keys())
            # pickle_module.dump(serialized_storage_keys, f, protocol=pickle_protocol)
            f.flush()

def load(f, map_location=None, pickle_module=pickle, **pickle_load_args):

    with _open_file_like(f, 'rb') as opened_file:
        pickle_file = 'data.pkl'
        if is_7zfile(opened_file):
            with _open_zipfile_reader(opened_file) as opened_zipfile:
                data_file = opened_zipfile.read(pickle_file)[pickle_file]
                unpickler = pickle_module.Unpickler(data_file, **pickle_load_args)
                # unpickler.persistent_load = persistent_load
                result = unpickler.load()
                return result
        else:
            f = opened_file
            magic_number = pickle_module.load(f, **pickle_load_args)
            if magic_number != MAGIC_NUMBER:
                raise RuntimeError("Invalid magic number; corrupt file?")
            protocol_version = pickle_module.load(f, **pickle_load_args)
            if protocol_version != PROTOCOL_VERSION:
                raise RuntimeError("Invalid protocol version: %s" % protocol_version)

            _sys_info = pickle_module.load(f, **pickle_load_args)
            unpickler = pickle_module.Unpickler(f, **pickle_load_args)
            # unpickler.persistent_load = persistent_load
            result = unpickler.load()
    return result


def _is_torchscript_zip(zip_file):
    return 'constants.pkl' in zip_file.get_all_records()

if __name__=="__main__":
    obj = {'k':'v'}
    save(obj,"model.pt")
    loaded = load("model.pt")
    print(loaded)