# Copyright (C) 2020-2021 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.

"""
Classes to host attributes from all layers.
"""

import copy
from enum import Enum
from typing import Dict, Optional, Union

from google.protobuf.message import Message
from here.content.content import Content, ContentBinding
from here.platform.adapter import Adapter
from here.platform.layer import VersionedLayer


class RangedAttribute:
    """
    A class to hold attributes with ranges
    """

    def __init__(self, start, end, attr):
        """
        Instantiate a RangedAttribute object.
        :param start: starting offset in [0, 1]
        :param end: ending offset in [0, 1], end >= start
        :param attr: an attribute of HMC, e.g. functional_class, speed_limit, etc.
        """
        self.start_offset = start
        self.end_offset = end
        self.attribute = attr

    def __str__(self):
        """
        Override the built-in str function.
        """
        return "start: {}, end: {}, attribute: {}".format(
            self.start_offset, self.end_offset, self.attribute
        )


class MultiSegmentAttribute:
    """
    A class to hold attributes with multiple segments
    """

    def __init__(self, start, end, oriented_segments, attr):
        """
        Instantiate a RangedAttribute object.
        :param start: starting offset in [0, 1]
        :param end: ending offset in [0, 1], end >= start
        :param oriented_segments: list of oriented segments
        :param attr: an attribute of HMC, e.g. functional_class, speed_limit, etc.
        """
        self.start_offset = start
        self.end_offset = end
        self.oriented_segments = oriented_segments
        self.attribute = attr

    def __str__(self):
        """
        Override the built-in str function.
        """
        return "start: {}, end: {},\noriented_segments: {},\nattribute: {}".format(
            self.start_offset, self.end_offset, self.oriented_segments, self.attribute
        )


def get_end_offset(anchor) -> float:
    """
    Take care of default value for last_segment_end_offset
    :param anchor: segment anchor
    :return: end_offset
    """
    end_offset = 1.0
    if anchor.HasField("last_segment_end_offset"):
        end_offset = anchor.last_segment_end_offset.value
    return end_offset


def enum_to_str(e: Union[Enum, str]) -> str:
    """Convert an enumerator or a string to a string"""
    return str(e.value) if isinstance(e, Enum) else e


class BaseLayer(ContentBinding):
    """
    A base class to hold general HERE Map Content (HMC) attributes
    as well as providing catalog, layer, and schema objects
    """

    def __init__(self, content: Content, layer_id: str, adapter: Optional[Adapter] = None):
        """
        Instantiate an HMC content layer.

        :param content: Content object
        :param layer_id: layer ID
        :param adapter: the ``Adapter`` to transform data between different representations,
            ``None`` to use the adapter defined in ``content``.
        """
        super().__init__(content, adapter)

        self.catalog = self.content._catalog
        self.layer = self.catalog.get_layer(layer_id)
        assert isinstance(self.layer, VersionedLayer)

        self.schema = self.layer.get_schema()

        # dictionary to hold attribute objects, per partition
        self.attributes: Dict[str, dict] = {}

    def get_partition(self, pid) -> Message:
        """
        Read a partition for given the partition ID

        :param pid: partition ID
        :return: the partition data, decoded into a Protobuf :class:`Message`
        :raises ValueError: if blob cannot be fetched and decoded because
            the partition contains no data or does not exist
        """

        # get the blob for the pid
        blob = next(
            self.layer.read_partitions([pid], decode=False, version=self.content.version), None
        )

        if blob is None:
            raise ValueError(
                f"Cannot fetch blob, the requested partition '{pid}' "
                "does not exist or contains no data."
            )

        data = blob[1]
        assert isinstance(data, bytes)

        # Possible improvement: the layer can decode the data directly
        decoded = self.schema.decode_blob(data)
        assert isinstance(decoded, Message)

        return decoded

    @staticmethod
    def resolve_single_segment_anchors(partition, attribute_name):
        """
        Associate single segment with attributes via segment_anchor
        :param partition: a partition object
        :param attribute_name: attribute collection name in the partition
        :return: a dict of segment -> attribute
        """

        attributes = {}

        for attr in getattr(partition, attribute_name):

            # clone this attribute and remove unnecessary fields for output
            value = copy.copy(attr)
            del value.segment_anchor_index[:]

            for idx in attr.segment_anchor_index:

                # get a segment anchor for this attribute
                anchor = partition.segment_anchor[idx]

                # get start offset for this anchor
                start_offset = anchor.first_segment_start_offset.value

                # get end offset for this anchor
                end_offset = get_end_offset(anchor)

                # 0: UNKNOWN, 1: BOTH, 2: FORWARD, 3: BACKWARD
                direction = anchor.attribute_orientation

                if "travel_direction" in attribute_name:
                    value = direction

                for osr in anchor.oriented_segment_ref:

                    # segmentId
                    sid = osr.segment_ref.identifier

                    # partitionId for this segment
                    pid = osr.segment_ref.partition_name

                    # initialize namap with this pid, sid, and direction
                    attributes.setdefault(pid, {})
                    attributes[pid].setdefault(sid, {})
                    attributes[pid][sid].setdefault(direction, [])
                    attributes[pid][sid][direction].append(
                        RangedAttribute(start_offset, end_offset, value)
                    )
                    attributes[pid][sid][direction].sort(key=lambda x: x.start_offset)

        return attributes

    @staticmethod
    def resolve_multi_segment_anchors(partition, attribute_name):
        """
        Associate multi-segments with attributes via segment_anchor
        :param partition: a partition object
        :param attribute_name: attribute collection name in the partition
        :return: a dict of segment -> attribute
        """

        attributes = {}

        for attr in getattr(partition, attribute_name):

            # clone this attribute and remove unnecessary fields for output
            value = copy.copy(attr)
            del value.segment_anchor_index[:]

            for idx in attr.segment_anchor_index:

                # get a segment anchor for this attribute
                anchor = partition.segment_anchor[idx]

                # get start offset for this anchor
                start_offset = anchor.first_segment_start_offset.value

                # get end offset for this anchor
                end_offset = get_end_offset(anchor)

                # 0: UNKNOWN, 1: BOTH, 2: FORWARD, 3: BACKWARD
                direction = anchor.attribute_orientation

                msa = MultiSegmentAttribute(
                    start_offset, end_offset, anchor.oriented_segment_ref, value
                )

                for osr in anchor.oriented_segment_ref:

                    # segmentId
                    sid = osr.segment_ref.identifier

                    # partitionId for this segment
                    pid = osr.segment_ref.partition_name

                    # initialize namap with this pid, sid, and direction
                    attributes.setdefault(pid, {})
                    attributes[pid].setdefault(sid, {})
                    attributes[pid][sid][direction] = msa

        return attributes

    def get_tabular_data_from_single_segment(self, pid, attribute_name):
        """
        Convert the attribute dict to tabular data
        :param pid: partition ID
        :param attribute_name: attribute name to look for
        :return: columns, data
        """

        columns = (
            "partitionId",
            "segmentId",
            "direction",
            "start_offset",
            "end_offset",
            "attribute",
        )
        data = []

        if attribute_name in self.attributes and pid in self.attributes[attribute_name]:
            for sid, directions in self.attributes[attribute_name][pid].items():
                for d, ras in directions.items():
                    for ra in ras:
                        data.append((pid, sid, d, ra.start_offset, ra.end_offset, ra.attribute))

        return columns, data

    def get_tabular_data_from_multi_segments(self, pid, attribute_name):
        """
        Convert the attribute dict to tabular data
        :param pid: partition ID
        :param attribute_name: attribute name to look for
        :return: columns, data
        """

        columns = (
            "partitionId",
            "segmentId",
            "direction",
            "start_offset",
            "end_offset",
            "oriented_segments",
            "attribute",
        )
        data = []

        if attribute_name in self.attributes and pid in self.attributes[attribute_name]:
            for sid, directions in self.attributes[attribute_name][pid].items():
                for d, msa in directions.items():
                    data.append(
                        (
                            pid,
                            sid,
                            d,
                            msa.start_offset,
                            msa.end_offset,
                            msa.oriented_segments,
                            msa.attribute,
                        )
                    )

        return columns, data
