# Copyright (C) 2020-2022 HERE Global B.V. and its affiliate(s).
# All rights reserved.
#
# This software and other materials contain proprietary information
# controlled by HERE and are protected by applicable copyright legislation.
# Any use and utilization of this software and other materials and
# disclosure to any third parties is conditional upon having a separate
# agreement with HERE for the access, use, utilization or disclosure of this
# software. In the absence of such agreement, the use of the software is not
# allowed.

"""
Protobuf Parser module
"""

import io
import json
import os
import zipfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

from google.protobuf.descriptor import FieldDescriptor
from google.protobuf.descriptor_pb2 import FileDescriptorSet
from google.protobuf.message import Message
from google.protobuf.message_factory import MessageFactory
from here.platform.exceptions import SchemaException

_LAYER_MANIFEST = "layer.manifest.json"
_LAYER_DESCRIPTORS = "layer.fds"
_PROTOBUF_EXTENSION = ".proto"


class ProtobufParser:
    """A class for parsing Protobuf data of a catalog layer partition."""

    def __init__(self, schema_pkg_file: Union[str, Path, io.BytesIO]):
        """
        Initialize a ProtobufParser instance.

        :param schema_pkg_file: the schema package zip-file as a str, Path or file-like object
        :raises SchemaException: if the required meta information is missing in schema package
        """
        layer_manifest, layer_descriptor_data = self._extract_layer_info(schema_pkg_file)
        if not layer_manifest:
            raise SchemaException(
                f"Invalid schema. Layer manifest file {_LAYER_MANIFEST} not found in schema "
                f"package file"
            )
        if not layer_descriptor_data:
            raise SchemaException(
                f"Invalid schema. Layer descriptor file {_LAYER_DESCRIPTORS} not found in schema "
                f"package file"
            )

        main_partition_class_name = json.loads(layer_manifest)["main"]["message"]

        descriptor_set = FileDescriptorSet.FromString(layer_descriptor_data)
        self.message_classes = self._create_message_classes(descriptor_set.file)

        if main_partition_class_name not in self.message_classes:
            raise SchemaException(
                 f"Unable to find main class {main_partition_class_name} present under" \
                 f"{_LAYER_MANIFEST} in {_LAYER_DESCRIPTORS} for schema artifact {schema_pkg_file}"
            )

        self.partition_class = self.message_classes[main_partition_class_name]

    @staticmethod
    def _extract_layer_info(
        input_zip: Union[str, Path, io.BytesIO]
    ) -> Tuple[Optional[str], Optional[bytes]]:
        """
        Extract layer info from an input zipfile.

        :param input_zip: input Protobuf schema zipfile
        :returns: content of the layer manifest and descriptor set
        """
        input_zipfile = zipfile.ZipFile(input_zip)
        layer_manifest = None
        layer_descriptor_set = None
        for name in input_zipfile.namelist():
            if name.endswith(_LAYER_MANIFEST):
                layer_manifest = input_zipfile.read(name).decode("utf-8")
            elif name.endswith(_LAYER_DESCRIPTORS):
                layer_descriptor_set = input_zipfile.read(name)

        return layer_manifest, layer_descriptor_set

    @staticmethod
    def _extract_protobuf_sources(input_zip: Union[str, Path, io.BytesIO]) -> Dict[str, str]:
        """
        Extract Protobuf from an input zipfile.

        :param input_zip: input Protobuf schema zipfile
        :returns: dictionary of schema file names and content
        """
        input_zipfile = zipfile.ZipFile(input_zip)
        protobuf_sources = {}
        for name in input_zipfile.namelist():
            if name.endswith(_PROTOBUF_EXTENSION):
                protobuf_sources[os.path.basename(name)] = input_zipfile.read(name).decode("utf-8")

        return protobuf_sources

    @staticmethod
    def _create_message_classes(file_protos) -> dict:
        """
        Create message classes.
        This code is taken from ``google.protobuf.message_factory.GetMessages``
        and modified to use a local message factory instance.

        :param file_protos: A list of protobuf files.
        :return: a dictionary mapping proto names to the message classes.
        """
        file_by_name = {file_proto.name: file_proto for file_proto in file_protos}

        factory = MessageFactory()

        def _add_file(file_proto):
            for dependency in file_proto.dependency:
                if dependency in file_by_name:
                    # Remove from elements to be visited, in order to cut cycles.
                    _add_file(file_by_name.pop(dependency))
            factory.pool.Add(file_proto)

        while file_by_name:
            _add_file(file_by_name.popitem()[1])

        return factory.GetMessages([file_proto.name for file_proto in file_protos])

    def parse_message(self, message: bytes) -> Message:
        """
        Parse a Protobuf message and returns the corresponding partition class instance.

        :param message: message received as blob
        :return: parsed Protobuf message
        """
        msg = self.partition_class.FromString(message)
        assert isinstance(msg, Message)
        return msg

    def primary_fields(self, all_fields: bool = False) -> List[FieldDescriptor]:
        """
        Return the field descriptors of the root partition class.

        :param all_fields: Return all field descriptors of the root partition
            class that are labeled 'repeated' if `all` is set to False.
            Return all the field descriptors otherwise.
        :return: List of field descriptors
        """
        return [
            f
            for f in self.partition_class.DESCRIPTOR.fields
            if all_fields or f.label == FieldDescriptor.LABEL_REPEATED
        ]

    @staticmethod
    def access_field(message: Message, path: Optional[str]):
        """
        Return the data at the given path.

        :param message: a decoded Protobuf message
        :param path: path to the field, concatenate field names with '.' to select nested fields
        :return: the field referenced by the ``path`` parameter. This is either
            a ``Message`` in case the field contains a nested Protobuf structure,
            a simple value for integers, strings and other native types,
            or a container that can be iterated for a repeated field in Protobuf.
        """
        if not path:
            return message
        else:
            current_field = message
            for field in path.split("."):
                current_field = getattr(current_field, field)
            return current_field

    def _access_message_class(self, class_name: str) -> Any:
        if class_name in self.message_classes:
            return self.message_classes[class_name]
        else:
            paths = class_name.split(".")
            for i in range(len(paths), 0, -1):
                current_parent_class = ".".join(paths[:i])
                if current_parent_class in self.message_classes:
                    result_class = self.message_classes[current_parent_class]
                    for path in paths[i:]:
                        result_class = result_class.DESCRIPTOR.nested_types_by_name[
                            path
                        ]._concrete_class
                    return result_class
