# 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, logging, queue, time
from typing import Dict, Tuple
from common.proto.context_pb2 import Connection, ConnectionId, Service, ServiceId, ServiceStatusEnum
from common.proto.pathcomp_pb2 import PathCompReply
from common.tools.grpc.Tools import grpc_message_to_json_string
from context.client.ContextClient import ContextClient
from service.service.service_handler_api.ServiceHandlerFactory import ServiceHandlerFactory
from service.service.tools.ObjectKeys import get_connection_key, get_service_key
from .tasks._Task import _Task
from .tasks.Task_ConnectionConfigure import Task_ConnectionConfigure
from .tasks.Task_ConnectionDeconfigure import Task_ConnectionDeconfigure
from .tasks.Task_ServiceDelete import Task_ServiceDelete
from .tasks.Task_ServiceSetStatus import Task_ServiceSetStatus
from .TaskExecutor import CacheableObjectType, TaskExecutor

LOGGER = logging.getLogger(__name__)

class TasksScheduler:
    def __init__(self, service_handler_factory : ServiceHandlerFactory) -> None:
        self._dag = graphlib.TopologicalSorter()
        self._executor = TaskExecutor(service_handler_factory)
        self._tasks : Dict[str, _Task] = dict()
        self._context_client = ContextClient()

    # ----- Helper methods ---------------------------------------------------------------------------------------------

    def _add_task_if_not_exists(self, task : _Task) -> str:
        task_key = task.key
        if task_key not in self._tasks:
            self._tasks[task_key] = task
        return task_key

    def _add_connection_to_executor_cache(self, connection : Connection) -> None:
        connection_key = get_connection_key(connection.connection_id)
        self._executor._store_editable_grpc_object(
            CacheableObjectType.CONNECTION, connection_key, Connection, connection)

    def _add_service_to_executor_cache(self, service : Service) -> None:
        service_key = get_service_key(service.service_id)
        self._executor._store_editable_grpc_object(
            CacheableObjectType.SERVICE, service_key, Service, service)

    # ----- Task & DAG composition methods -----------------------------------------------------------------------------

    def _service_create(self, service_id : ServiceId) -> Tuple[str, str]:
        service_planned_key = self._add_task_if_not_exists(Task_ServiceSetStatus(
            self._executor, service_id, ServiceStatusEnum.SERVICESTATUS_PLANNED))

        service_active_key = self._add_task_if_not_exists(Task_ServiceSetStatus(
            self._executor, service_id, ServiceStatusEnum.SERVICESTATUS_ACTIVE))

        # activating a service requires the service is in planning state
        self._dag.add(service_active_key, service_planned_key)
        return service_planned_key, service_active_key

    def _service_remove(self, service_id : ServiceId) -> Tuple[str, str]:
        service_removing_key = self._add_task_if_not_exists(Task_ServiceSetStatus(
            self._executor, service_id, ServiceStatusEnum.SERVICESTATUS_PENDING_REMOVAL))

        service_delete_key = self._add_task_if_not_exists(Task_ServiceDelete(self._executor, service_id))

        # deleting a service requires the service is in removing state
        self._dag.add(service_delete_key, service_removing_key)
        return service_removing_key, service_delete_key

    def _connection_configure(self, connection_id : ConnectionId, service_id : ServiceId) -> str:
        connection_configure_key = self._add_task_if_not_exists(Task_ConnectionConfigure(
            self._executor, connection_id))

        # the connection configuration depends on its connection's service being in planning state
        service_planned_key = self._add_task_if_not_exists(Task_ServiceSetStatus(
            self._executor, service_id, ServiceStatusEnum.SERVICESTATUS_PLANNED))
        self._dag.add(connection_configure_key, service_planned_key)

        # the connection's service depends on the connection configuration to transition to active state
        service_active_key = self._add_task_if_not_exists(Task_ServiceSetStatus(
            self._executor, service_id, ServiceStatusEnum.SERVICESTATUS_ACTIVE))
        self._dag.add(service_active_key, connection_configure_key)

        return connection_configure_key

    def _connection_deconfigure(self, connection_id : ConnectionId, service_id : ServiceId) -> str:
        connection_deconfigure_key = self._add_task_if_not_exists(Task_ConnectionDeconfigure(
            self._executor, connection_id))

        # the connection deconfiguration depends on its connection's service being in removing state
        service_pending_removal_key = self._add_task_if_not_exists(Task_ServiceSetStatus(
            self._executor, service_id, ServiceStatusEnum.SERVICESTATUS_PENDING_REMOVAL))
        self._dag.add(connection_deconfigure_key, service_pending_removal_key)

        # the connection's service depends on the connection deconfiguration to transition to delete
        service_delete_key = self._add_task_if_not_exists(Task_ServiceDelete(
            self._executor, service_id))
        self._dag.add(service_delete_key, connection_deconfigure_key)

        return connection_deconfigure_key

    def compose_from_pathcompreply(self, pathcomp_reply : PathCompReply, is_delete : bool = False) -> None:
        t0 = time.time()
        include_service = self._service_remove if is_delete else self._service_create
        include_connection = self._connection_deconfigure if is_delete else self._connection_configure

        for service in pathcomp_reply.services:
            include_service(service.service_id)
            self._add_service_to_executor_cache(service)

        for connection in pathcomp_reply.connections:
            connection_key = include_connection(connection.connection_id, connection.service_id)
            self._add_connection_to_executor_cache(connection)
            self._executor.get_service(connection.service_id)
            for sub_service_id in connection.sub_service_ids:
                _,service_key_done = include_service(sub_service_id)
                self._executor.get_service(sub_service_id)
                self._dag.add(connection_key, service_key_done)

        t1 = time.time()
        LOGGER.info('[compose_from_pathcompreply] elapsed_time: {:f} sec'.format(t1-t0))

    def compose_from_service(self, service : Service, is_delete : bool = False) -> None:
        t0 = time.time()
        include_service = self._service_remove if is_delete else self._service_create
        include_connection = self._connection_deconfigure if is_delete else self._connection_configure

        explored_items = set()
        pending_items_to_explore = queue.Queue()
        pending_items_to_explore.put(service)

        while not pending_items_to_explore.empty():
            try:
                item = pending_items_to_explore.get(block=False)
            except queue.Empty:
                break

            if isinstance(item, Service):
                str_item_key = grpc_message_to_json_string(item.service_id)
                LOGGER.info('Exploring Service: {:s}'.format(str_item_key))
                if str_item_key in explored_items: continue

                include_service(item.service_id)
                self._add_service_to_executor_cache(item)
                connections = self._context_client.ListConnections(item.service_id)
                for connection in connections.connections:
                    self._add_connection_to_executor_cache(connection)
                    LOGGER.info('  PUT connection {:s}'.format(grpc_message_to_json_string(connection.connection_id)))
                    pending_items_to_explore.put(connection)

                explored_items.add(str_item_key)

            elif isinstance(item, ServiceId):
                str_item_key = grpc_message_to_json_string(item)
                LOGGER.info('Exploring ServiceId: {:s}'.format(str_item_key))
                if str_item_key in explored_items: continue

                include_service(item)
                self._executor.get_service(item)
                connections = self._context_client.ListConnections(item)
                for connection in connections.connections:
                    self._add_connection_to_executor_cache(connection)
                    LOGGER.info('  PUT connection {:s}'.format(grpc_message_to_json_string(connection.connection_id)))
                    pending_items_to_explore.put(connection)

                explored_items.add(str_item_key)

            elif isinstance(item, Connection):
                str_item_key = grpc_message_to_json_string(item.connection_id)
                LOGGER.info('Exploring Connection: {:s}'.format(str_item_key))
                if str_item_key in explored_items: continue

                connection_key = include_connection(item.connection_id, item.service_id)
                self._add_connection_to_executor_cache(connection)

                #_,service_key_done = include_service(item.service_id)
                self._executor.get_service(item.service_id)
                #self._dag.add(service_key_done, connection_key)
                LOGGER.info('  PUT service_id {:s}'.format(grpc_message_to_json_string(item.service_id)))
                pending_items_to_explore.put(item.service_id)

                for sub_service_id in item.sub_service_ids:
                    _,service_key_done = include_service(sub_service_id)
                    self._executor.get_service(sub_service_id)
                    self._dag.add(service_key_done, connection_key)
                    LOGGER.info('  PUT sub_service_id {:s}'.format(grpc_message_to_json_string(sub_service_id)))
                    pending_items_to_explore.put(sub_service_id)

                explored_items.add(str_item_key)

            else:
                MSG = 'Unsupported item {:s}({:s})'
                raise Exception(MSG.format(type(item).__name__, grpc_message_to_json_string(item)))

        t1 = time.time()
        LOGGER.info('explored_items={:s}'.format(str(explored_items)))
        LOGGER.info('[compose_from_service] elapsed_time: {:f} sec'.format(t1-t0))

    def execute_all(self, dry_run : bool = False) -> None:
        ordered_task_keys = list(self._dag.static_order())
        LOGGER.info('ordered_task_keys={:s}'.format(str(ordered_task_keys)))

        results = []
        for task_key in ordered_task_keys:
            LOGGER.info('Task {:s} - begin'.format(str(task_key)))
            task = self._tasks.get(task_key)
            succeeded = True if dry_run else task.execute()
            LOGGER.info('Task {:s} - succeeded={:s}'.format(str(task_key), str(succeeded)))
            results.append(succeeded)

        LOGGER.info('execute_all results={:s}'.format(str(results)))
        return zip(ordered_task_keys, results)
