from manifest_edit.plugin.mpd import ManifestIteratorPlugin
from manifest_edit.context import Context
from manifest_edit import libfmp4
from schema import Schema, Optional, Or
from manifest_edit.plugins.utility import renderFormatString
import re


class Plugin(ManifestIteratorPlugin):
    """
    Segment URL modifier plugin

    This plugin allows to modify the 'initialization' and 'media' attributes of
    the SegmentTemplate element (see 5.3.9). This is an example of such an
    element:

    .. code-block:: xml

      <SegmentTemplate
        timescale="48000"
        initialization="audio-$RepresentationID$.dash"
        media="audio-$RepresentationID$-$Time$.dash">
        <SegmentTimeline>
          <S t="0" d="125952" />
          <S d="128000" r="272" />
          <S d="158992" />
        </SegmentTimeline>
      </SegmentTemplate>

    It also allows to modify the baseURL element wherever it may be
    present in the manifest.

    In particular scenarios where users may want to serve segments from
    different servers or referring to individual ism server manifests,
    editing these attributes is essential to form the desired segment URLS.
    """

    _name = __name__

    _keys = ["segmentTemplate", "baseURL"]

    # This is just the specific ["how" config] for this plugin
    def schema(self):
        return Schema(
            {
                Optional(self._keys[0], default={}): {
                    Optional("initialization", default={}): Or(
                        {},
                        {
                            "match": lambda s: re.compile(s),
                            "replace": str,
                        },
                    ),
                    Optional("media", default={}): Or(
                        {},
                        {
                            "match": lambda s: re.compile(s),
                            "replace": str,
                        },
                    ),
                },
                Optional(self._keys[1], default={}): Or(
                    {},
                    {
                        "match": lambda s: re.compile(s),
                        "replace": str,
                    },
                    {
                        "URL": Or(
                            {},
                            {
                                "match": lambda s: re.compile(s),
                                "replace": str,
                            },
                        ),
                    },
                    {
                        "serviceLocation": Or(
                            {},
                            {
                                "match": lambda s: re.compile(s),
                                "replace": str,
                            },
                        ),
                    }
                ),
            }
        )

    def _attemptMod(self, value: str, match: str, replace: str) -> str:
        # We know that "match" is a proper regular expression
        # We don't know if the replace string is. If it is not, re.sub will
        # just return the original string

        (new_string, number_of_subs_made) = re.subn(match, replace, value)

        if number_of_subs_made == 0:
            # This is actually OK, it is possible that a user has selected
            # all representations (i.e. - id: '.*') but then chose a match
            # string that only matches some representations.
            Context.log_trace(
                f"When trying to modify string '{value}' "
                f"using match string '{match}' "
                f"and replace string '{replace}' "
                " no substitutions have been made!"
            )

        return new_string

    def _modifySegmentTemplate(self, segment_config, element):
        for attribute_name in ["initialization", "media"]:
            if segment_config[attribute_name]:
                # segmentTemplate is an optional attribute and as such
                # you must check if it's there. If it's not, just skip...
                if getattr(element, "segmentTemplate", None):
                    rendered_replace_string = renderFormatString(
                        element, segment_config[attribute_name]["replace"]
                    )
                    # Weird stuff can happen with supposedly "raw" string
                    # coming from pyyaml and unicode escapes, so prepare to
                    # debug
                    Context.log_trace(
                        f"Rendered string '{segment_config[attribute_name]['replace']}' into '{rendered_replace_string}'"
                    )
                    setattr(
                        element.segmentTemplate,
                        attribute_name,
                        self._attemptMod(
                            getattr(element.segmentTemplate, attribute_name),
                            segment_config[attribute_name]["match"],
                            rendered_replace_string,
                        ),
                    )
                else:
                    # We will not allow to *create* a segment template from
                    # scratch in manifest edit because of course this will
                    # never work. 
                    Context.log_trace(
                        f'No segmentTemplate element detected in this '
                        f'{type(element)}. Is the configuration correct?'
                    )

    def _modifyBaseURLByURL(self, baseURL_config, element):
        if baseURL_config:
            rendered_replace_string = renderFormatString(
                element, baseURL_config["replace"]
            )
            Context.log_trace(
                f"Rendered string '{baseURL_config['replace']}' into '{rendered_replace_string}'"
            )
            # CHANGE RELATED TO INTRODUCTION OF MULTIPLE BASEURLs.
            # element.baseURLs is a list that can be empty. 
            # Previously, when only baseURL was present, it defaulted to an
            # empty string. For this reason, it was expected that when a
            # user-provided regular expression matches an empty string,
            # the _attemptMod function creates a baseURL entry.
            # This still works for non-empty baseURLs list, otherwise you'll
            # have to append to it.
            if not len(element.baseURLs):
                # Check if user-provided regex matches the empty string
                new_url = self._attemptMod(
                        "",
                        baseURL_config["match"],
                        rendered_replace_string,
                    )

                if new_url:
                    element.baseURLs.append(libfmp4.mpd.BaseURL(new_url))

            else:
                for baseURL in element.baseURLs:
                    # baseURL needs to be converted back and forth to string
                    setattr(
                        baseURL,
                        "URL",
                        libfmp4.Url(
                            self._attemptMod(
                                str(baseURL.URL),
                                baseURL_config["match"],
                                rendered_replace_string,
                            )
                        ),
                    )

    def _modifyBaseURLByServiceLocation(self, baseURL_config, element):
        if baseURL_config:
            rendered_replace_string = renderFormatString(
                element, baseURL_config["replace"]
            )
            Context.log_trace(
                f"Rendered string '{baseURL_config['replace']}' into '{rendered_replace_string}'"
            )
            # When it comes to serviceLocation, we are only going to allow changing existing
            # entries that have an URL attribute.
            if not len(element.baseURLs):
                Context.log_error(
                    "Attempt to replace/assign baseURL.serviceLocation on an element with"
                    " no baseURL at all. This is not possible: instead you should"
                    " first create a valid baseURL entry with an URL attribute and"
                    " empty serviceLocation, then you can change serviceLocation with"
                    " a second pass of this plugin."
                )

            else:
                for baseURL in element.baseURLs:
                    # serviceLocation needs to be converted back and forth to string
                    # Annoyingly enough, serviceLocation defaults to None, as
                    # opposed to baseURL.URL that defaults to empty string
                    # You don't want to turn None into an empty string with
                    # _attemptMpd
                    if baseURL.serviceLocation is not None:
                        # treat as a string, possibly empty
                        setattr(
                            baseURL,
                            "serviceLocation",
                            self._attemptMod(
                                str(baseURL.serviceLocation),
                                baseURL_config["match"],
                                rendered_replace_string,
                            )
                        )
                    else:
                        new_val = self._attemptMod(
                            '',
                            baseURL_config["match"],
                            rendered_replace_string,
                        )
                        if new_val:
                            setattr(
                                baseURL,
                                "serviceLocation",
                                new_val
                            )

    def _modifyBaseURL(self, baseURL_config, element):
        # if "match" or "release" are present, fallback to old behaviour
        # of defaulting to use baseURL.URL attribute
        # Otherwise check the specific attribute the user wants to 
        # search/replace, either URL or serviceLocation
        if any([key in baseURL_config.keys() for key in ["match", "replace"]]):
            self._modifyBaseURLByURL(baseURL_config, element)
        elif "URL" in baseURL_config.keys():
            self._modifyBaseURLByURL(baseURL_config["URL"], element)
        elif "serviceLocation" in baseURL_config.keys():
            self._modifyBaseURLByServiceLocation(baseURL_config["serviceLocation"], element)

    def _modifyBaseURLOrSegmentTemplate(self, manifest, storage):
        for segment_config, element in self.config(manifest, storage):
            if segment_config.get(self._keys[1], None):
                self._modifyBaseURL(segment_config[self._keys[1]], element)
            elif segment_config.get(self._keys[0], None):
                self._modifySegmentTemplate(segment_config[self._keys[0]], element)

    def process(self, manifest, storage):
        self._modifyBaseURLOrSegmentTemplate(manifest, storage)
