from UM.Workspace.WorkspaceWriter import WorkspaceWriter
from UM.Application import Application
from UM.Preferences import Preferences
from UM.Settings.ContainerRegistry import ContainerRegistry
from UM.Settings.ContainerStack import ContainerStack
from cura.Settings.ExtruderManager import ExtruderManager
import zipfile
from io import StringIO
import copy
import configparser


class ThreeMFWorkspaceWriter(WorkspaceWriter):
    def __init__(self):
        super().__init__()

    def write(self, stream, nodes, mode=WorkspaceWriter.OutputMode.BinaryMode):
        mesh_writer = Application.getInstance().getMeshFileHandler().getWriter("3MFWriter")

        if not mesh_writer:  # We need to have the 3mf mesh writer, otherwise we can't save the entire workspace
            return False

        # Indicate that the 3mf mesh writer should not close the archive just yet (we still need to add stuff to it).
        mesh_writer.setStoreArchive(True)
        mesh_writer.write(stream, nodes, mode)

        archive = mesh_writer.getArchive()
        if archive is None:  # This happens if there was no mesh data to write.
            archive = zipfile.ZipFile(stream, "w", compression = zipfile.ZIP_DEFLATED)

        global_container_stack = Application.getInstance().getGlobalContainerStack()

        # Add global container stack data to the archive.
        self._writeContainerToArchive(global_container_stack, archive)

        # Also write all containers in the stack to the file
        for container in global_container_stack.getContainers():
            self._writeContainerToArchive(container, archive)

        # Check if the machine has extruders and save all that data as well.
        for extruder_stack in ExtruderManager.getInstance().getMachineExtruders(global_container_stack.getId()):
            self._writeContainerToArchive(extruder_stack, archive)
            for container in extruder_stack.getContainers():
                self._writeContainerToArchive(container, archive)

        # Write preferences to archive
        preferences_file = zipfile.ZipInfo("Cura/preferences.cfg")
        preferences_string = StringIO()
        Preferences.getInstance().writeToFile(preferences_string)
        archive.writestr(preferences_file, preferences_string.getvalue())

        # Save Cura version
        version_file = zipfile.ZipInfo("Cura/version.ini")
        version_config_parser = configparser.ConfigParser()
        version_config_parser.add_section("versions")
        version_config_parser.set("versions", "cura_version", Application.getStaticVersion())

        version_file_string = StringIO()
        version_config_parser.write(version_file_string)
        archive.writestr(version_file, version_file_string.getvalue())

        # Close the archive & reset states.
        archive.close()
        mesh_writer.setStoreArchive(False)
        return True

    ##  Helper function that writes ContainerStacks, InstanceContainers and DefinitionContainers to the archive.
    #   \param container That follows the \type{ContainerInterface} to archive.
    #   \param archive The archive to write to.
    @staticmethod
    def _writeContainerToArchive(container, archive):
        if type(container) == type(ContainerRegistry.getInstance().getEmptyInstanceContainer()):
            return  # Empty file, do nothing.

        file_suffix = ContainerRegistry.getMimeTypeForContainer(type(container)).preferredSuffix

        # Some containers have a base file, which should then be the file to use.
        if "base_file" in container.getMetaData():
            base_file = container.getMetaDataEntry("base_file")
            container = ContainerRegistry.getInstance().findContainers(id = base_file)[0]

        file_name = "Cura/%s.%s" % (container.getId(), file_suffix)

        if file_name in archive.namelist():
            return  # File was already saved, no need to do it again. Uranium guarantees unique ID's, so this should hold.

        file_in_archive = zipfile.ZipInfo(file_name)
        # For some reason we have to set the compress type of each file as well (it doesn't keep the type of the entire archive)
        file_in_archive.compress_type = zipfile.ZIP_DEFLATED
        if type(container) == ContainerStack and (container.getMetaDataEntry("network_authentication_id") or container.getMetaDataEntry("network_authentication_key")):
            # TODO: Hack
            # Create a shallow copy of the container, so we can filter out the network auth (if any)
            container_copy = copy.deepcopy(container)
            container_copy.removeMetaDataEntry("network_authentication_id")
            container_copy.removeMetaDataEntry("network_authentication_key")
            serialized_data = container_copy.serialize()
        else:
            serialized_data = container.serialize()

        archive.writestr(file_in_archive, serialized_data)