from collective.ftw.upgrade import UpgradeStep
from plone import api

import logging
import re

logger = logging.getLogger("__name__")


class BogusPattern:
    def __init__(self, pattern, exception):
        self.pattern = pattern
        self.exception = exception

    def __repr__(self):
        return f"BogusPattern({self.pattern!r}, {self.exception!r})"


class TemporaryReCompile:
    """See also https://github.com/plone/plone.app.upgrade/pull/334"""

    original_compile = re._compile

    @staticmethod
    def custom_compile(pattern, flags=0):
        """
        Custom compile function to handle problematic patterns while unpickling.
        """
        try:
            return TemporaryReCompile.original_compile(pattern, flags)
        except re.error as exception:
            return BogusPattern(pattern, exception)

    def __enter__(self):
        """
        Replace `re.compile` with a custom compile function that
        will return a BogusPattern instance in case the pattern cannot be compiled.
        """
        re._compile = TemporaryReCompile.custom_compile

    def __exit__(self, exc_type, exc_value, traceback):
        """
        Restore the original `re.compile` function.
        """
        re._compile = TemporaryReCompile.original_compile


class UpgradeTheMimetypeRegistryEntries(UpgradeStep):
    """Upgrade the mimetype registry entries."""

    def __call__(self):
        mtr = api.portal.get_tool("mimetypes_registry")
        bad_globs = {}

        logger.info("Checking the mimetypes_registry globs")

        # We use a context manager to provide a custom `re.compile` function
        # that will allow to unpickled globs with patterns that are not compatible
        # with the current Python version.
        with TemporaryReCompile():
            for glob in set(mtr.globs):
                pattern, mimetype = mtr.globs[glob]
                if isinstance(pattern, BogusPattern):
                    logger.info(
                        "Found problematic glob %r with pattern %r", glob, pattern
                    )
                    bad_globs[glob] = mimetype

            # While in the context manager we can safely remove the problematic globs
            for glob in bad_globs:
                del mtr.globs[glob]

        # Re-add the problematic globs with the correct pattern
        for glob, mimetype in bad_globs.items():
            mtr.register_glob(glob, mimetype)

        logger.info("%r globs were fixed", len(bad_globs))
