# Copyright 2021-2023 H2020 TeraFlow (https://www.teraflow-h2020.eu/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import graphlib
from enum import Enum
from typing import Dict, List, Optional, Tuple, Union
from common.proto.context_pb2 import Connection, ConnectionId, Service, ServiceId
from common.proto.pathcomp_pb2 import PathCompReply

# Compose Directed Acyclic Graph of dependencies between connections and services
# retrieved by PathComp to create them in the appropriate order.

class ObjectType(Enum):
    CONNECTION = 'connection'
    SERVICE    = 'service'

ObjectKey  = Tuple[ObjectType, str]
ObjectId   = Union[ConnectionId, ServiceId]
ObjectData = Union[Connection, Service]
ObjectItem = Tuple[ObjectId, Optional[ObjectData]]
ObjectDict = Dict[ObjectKey, ObjectItem]
Resolution = List[Tuple[ObjectKey, ObjectItem]]

def get_connection_key(connection_id : ConnectionId) -> ObjectKey:
    connection_uuid = connection_id.connection_uuid.uuid
    return ObjectType.CONNECTION.value, connection_uuid

def get_service_key(service_id : ServiceId) -> ObjectKey:
    context_uuid = service_id.context_id.context_uuid.uuid
    service_uuid = service_id.service_uuid.uuid
    return ObjectType.SERVICE.value, '/'.join([context_uuid, service_uuid])

def resolve_dependencies(pathcomp_reply : PathCompReply) -> Resolution:
    dag = graphlib.TopologicalSorter()
    objects : ObjectDict = dict()

    for service in pathcomp_reply.services:
        service_key = get_service_key(service.service_id)
        objects[service_key] = (service.service_id, service)

    for connection in pathcomp_reply.connections:
        connection_key = get_connection_key(connection.connection_id)
        objects[connection_key] = (connection.connection_id, connection)

        # the connection's service depends on the connection
        service_key = get_service_key(connection.service_id)
        dag.add(service_key, connection_key)
        if service_key not in objects: objects[service_key] = (connection.service_id, None)

        # the connection depends on these sub-services
        for sub_service_id in connection.sub_service_ids:
            sub_service_key = get_service_key(sub_service_id)
            dag.add(connection_key, sub_service_key)
            if sub_service_key not in objects: objects[sub_service_key] = (sub_service_id, None)

    resolution : Resolution = list()
    for item_key in dag.static_order():
        item_tuple = objects.get(item_key)
        resolution.append((item_key, item_tuple))

    return resolution
