# Copyright (C) 2020-2023 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 os
import zipfile
from pathlib import Path
from typing import Any, Dict, List, Optional, 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
from here.platform.schema.parser import Parser

_PROTOBUF_EXTENSION = ".proto"


class ProtobufParser(Parser):
    """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 = Parser.extract_layer_info(schema_pkg_file)
        if not layer_manifest:
            raise SchemaException(
                f"Invalid schema. Layer manifest file {Parser.LAYER_MANIFEST} not found in schema "
                "package file"
            )
        if not layer_descriptor_data:
            raise SchemaException(
                f"Invalid schema. Layer descriptor file {Parser.LAYER_DESCRIPTORS} "
                "not found in schema package file"
            )

        main_partition_class_name = layer_manifest["main"]["message"]

        descriptor_set: FileDescriptorSet = FileDescriptorSet.FromString(layer_descriptor_data)
        self.message_classes = ProtobufParser._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"{Parser.LAYER_MANIFEST} in {Parser.LAYER_DESCRIPTORS} for the schema artifact"
            )

        self.partition_class = self.message_classes[main_partition_class_name]

    @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 get_partition_class(self):
        """
        Returns partition class.

        :return: Partition class
        """
        return self.partition_class

    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 parse_message_by_class(
        self, message: bytes, fully_qualified_message_name: str
    ) -> Union[Message, Any]:
        """
        Parse a Protobuf message by proto class and returns the corresponding
        partition class instance.

        :param message: message received as blob
        :param fully_qualified_message_name: fully qualified message name of the proto class
        :return: parsed Protobuf message
        """
        fully_qualified_message_name = fully_qualified_message_name.strip()
        if fully_qualified_message_name in self.message_classes:
            msg = self.message_classes[fully_qualified_message_name].FromString(message)
            assert isinstance(msg, Message)
            return msg
        return None

    def get_message_class(self, fully_qualified_message_name: str) -> Any:
        """
        Get Message class by fully qualified message name.

        :param fully_qualified_message_name: fully qualified message name of the proto class
        :return: message class
        """
        fully_qualified_message_name = fully_qualified_message_name.strip()
        if fully_qualified_message_name in self.message_classes:
            return self.message_classes[fully_qualified_message_name]
        return None

    def serialize_bytes(self, data: Union[Message, dict]) -> bytes:
        """
        Parse a Protobuf message and returns the bytes.

        :param data: Protobuf message
        :return: bytes
        """
        assert isinstance(data, Message)
        message_bytes = data.SerializeToString()
        assert isinstance(message_bytes, bytes)
        return message_bytes

    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
