# 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 TYPE_CHECKING, Dict, Union

from here.platform.layer import VersionedLayer

if TYPE_CHECKING:
    from here.content.hmc.hmc import HMC


class RangedAttribute(object):
    """
    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(object):
    """
    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(object):
    """
    A base class to hold general HERE Map Content (HMC) attributes
    as well as providing catalog, layer, and schema objects
    """

    def __init__(self, content: "HMC", layer_id: str):
        """
        Instantiate an HMC content layer.

        :param content: instance of the HMC root class
        :param layer_id: layer ID
        """
        # content, catalog and adapter objects for HMC
        self.content = content
        self.adapter = content.adapter

        self.layer: VersionedLayer = content.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):
        """
        Read a partition for given the partition ID

        :param pid: partition ID
        :return: a partition object, decoded into a dictionary
        :raises ValueError: if blob is not properly fetched
        """

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

        if blob is None:
            raise ValueError("blob is not properly fetched.")

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

        # TODO: ask the layer to decode directly, instead of doing it here manually
        decoded = self.schema.decode_blob(data)
        # assert isinstance(decoded, dict)

        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
        """

        data = []
        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 (
            ("partitionId", "segmentId", "direction", "start_offset", "end_offset", "attribute"),
            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
        """

        data = []
        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 (
            (
                "partitionId",
                "segmentId",
                "direction",
                "start_offset",
                "end_offset",
                "oriented_segments",
                "attribute",
            ),
            data,
        )
