# Copyright 2022-2024 ETSI OSG/SDG TeraFlowSDN (TFS) (https://tfs.etsi.org/)
#
# 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 json, logging
from enum import Enum
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
from common.DeviceTypes import DeviceTypeEnum
from common.method_wrappers.ServiceExceptions import NotFoundException
from typing import List
from common.proto.qkd_app_pb2 import QKDAppStatusEnum
from common.proto.context_pb2 import (
    Connection, ConnectionId, Device, DeviceDriverEnum, DeviceId, Service, ServiceId,
    OpticalConfig, OpticalConfigId
)
from common.proto.qkd_app_pb2 import App, AppId
from common.proto.context_pb2 import ContextId
from common.tools.context_queries.Connection import get_connection_by_id
from common.tools.context_queries.Device import get_device
from common.tools.context_queries.Service import get_service_by_id
from common.tools.grpc.Tools import grpc_message_to_json_string
from common.tools.object_factory.Device import json_device_id
from context.client.ContextClient import ContextClient
from device.client.DeviceClient import DeviceClient
from qkd_app.client.QKDAppClient import QKDAppClient
from service.service.service_handler_api.Exceptions import (
    UnsatisfiedFilterException, UnsupportedFilterFieldException, UnsupportedFilterFieldValueException
)
from service.service.service_handler_api.ServiceHandlerFactory import ServiceHandlerFactory, get_service_handler_class
from service.service.tools.ObjectKeys import get_connection_key, get_device_key, get_service_key, get_qkd_app_key

if TYPE_CHECKING:
    from service.service.service_handler_api._ServiceHandler import _ServiceHandler

LOGGER = logging.getLogger(__name__)

CacheableObject = Union[Connection, Device, Service]

class CacheableObjectType(Enum):
    CONNECTION = 'connection'
    DEVICE     = 'device'
    SERVICE    = 'service'
    QKD_APP    = 'qkd-app'

class TaskExecutor:
    def __init__(self, service_handler_factory : ServiceHandlerFactory) -> None:
        self._service_handler_factory = service_handler_factory
        self._context_client = ContextClient()
        # DEPENDENCY QKD
        self._qkd_app_client = QKDAppClient()
        self._device_client = DeviceClient()
        self._grpc_objects_cache : Dict[str, CacheableObject] = dict()

    @property
    def service_handler_factory(self) -> ServiceHandlerFactory: return self._service_handler_factory

    # ----- Common methods ---------------------------------------------------------------------------------------------

    def _load_grpc_object(self, object_type : CacheableObjectType, object_key : str) -> Optional[CacheableObject]:
        object_key = '{:s}:{:s}'.format(object_type.value, object_key)
        return self._grpc_objects_cache.get(object_key)

    def _store_grpc_object(self, object_type : CacheableObjectType, object_key : str, grpc_object) -> None:
        object_key = '{:s}:{:s}'.format(object_type.value, object_key)
        self._grpc_objects_cache[object_key] = grpc_object
    
    def _delete_grpc_object(self, object_type : CacheableObjectType, object_key : str) -> None:
        object_key = '{:s}:{:s}'.format(object_type.value, object_key)
        self._grpc_objects_cache.pop(object_key, None)

    def _store_editable_grpc_object(
        self, object_type : CacheableObjectType, object_key : str, grpc_class, grpc_ro_object
    ) -> Any:
        grpc_rw_object = grpc_class()
        grpc_rw_object.CopyFrom(grpc_ro_object)
        self._store_grpc_object(object_type, object_key, grpc_rw_object)
        return grpc_rw_object

    # ----- Connection-related methods ---------------------------------------------------------------------------------

    def get_connection(self, connection_id : ConnectionId) -> Connection:
        connection_key = get_connection_key(connection_id)
        connection = self._load_grpc_object(CacheableObjectType.CONNECTION, connection_key)
        if connection is None:
            connection = get_connection_by_id(self._context_client, connection_id)
            if connection is None: raise NotFoundException('Connection', connection_key)
            connection : Connection = self._store_editable_grpc_object(
                CacheableObjectType.CONNECTION, connection_key, Connection, connection)
        return connection

    def set_connection(self, connection : Connection) -> None:
        connection_key = get_connection_key(connection.connection_id)
        self._context_client.SetConnection(connection)
        self._store_grpc_object(CacheableObjectType.CONNECTION, connection_key, connection)

    def delete_connection(self, connection_id : ConnectionId) -> None:
        connection_key = get_connection_key(connection_id)
        self._context_client.RemoveConnection(connection_id)
        self._delete_grpc_object(CacheableObjectType.CONNECTION, connection_key)

    # ----- Device-related methods -------------------------------------------------------------------------------------

    def get_device(self, device_id : DeviceId) -> Device:
        device_key = get_device_key(device_id)
        device = self._load_grpc_object(CacheableObjectType.DEVICE, device_key)
        if device is None:
            device = get_device(self._context_client, device_id.device_uuid.uuid)
            if device is None: raise NotFoundException('Device', device_key)
            device : Device = self._store_editable_grpc_object(
                CacheableObjectType.DEVICE, device_key, Device, device)
        return device

    def configure_device(self, device : Device) -> None:
        device_key = get_device_key(device.device_id)
        self._device_client.ConfigureDevice(device)
        self._store_grpc_object(CacheableObjectType.DEVICE, device_key, device)
    
    # New function Andrea for Optical Devices
    def configure_optical_device(self, device : Device, settings : str, flows : list, is_opticalband : bool):
        device_key = get_device_key(device.device_id)
        optical_config_id = OpticalConfigId()
        optical_config_id.opticalconfig_uuid = device.device_id.device_uuid.uuid
        optical_config = OpticalConfig()
        setting = settings.value if settings else ''

        try:
            result = self._context_client.SelectOpticalConfig(optical_config_id)
            if result is not None:
                new_config = json.loads(result.config)
                new_config["new_config"] = setting
                new_config["is_opticalband"] = is_opticalband
                new_config["flow"] = flows
                result.config = json.dumps(new_config)
                optical_config.CopyFrom(result)
                self._device_client.ConfigureOpticalDevice(optical_config)
            self._store_grpc_object(CacheableObjectType.DEVICE, device_key, device)
        except Exception as e:
            LOGGER.info("error in config my config %s",e)

    def get_device_controller(self, device : Device) -> Optional[Device]:
        #json_controller = None
        #for config_rule in device.device_config.config_rules:
        #    if config_rule.WhichOneof('config_rule') != 'custom': continue
        #    if config_rule.custom.resource_key != '_controller': continue
        #    json_controller = json.loads(config_rule.custom.resource_value)
        #    break

        #if json_controller is None: return None

        #controller_uuid = json_controller['uuid']
        controller_uuid = device.controller_id.device_uuid.uuid
        if len(controller_uuid) == 0: return None
        controller = self.get_device(DeviceId(**json_device_id(controller_uuid)))
        controller_uuid = controller.device_id.device_uuid.uuid
        if controller is None: raise Exception('Device({:s}) not found'.format(str(controller_uuid)))
        return controller

    def get_devices_from_connection(
        self, connection : Connection, exclude_managed_by_controller : bool = False
    ) -> Dict[DeviceTypeEnum, Dict[str, Device]]:
        devices : Dict[DeviceTypeEnum, Dict[str, Device]] = dict()
        for endpoint_id in connection.path_hops_endpoint_ids:
            device = self.get_device(endpoint_id.device_id)
            device_uuid = endpoint_id.device_id.device_uuid.uuid
            if device is None: raise Exception('Device({:s}) not found'.format(str(device_uuid)))

            controller = self.get_device_controller(device)
            if controller is None:
                device_type = DeviceTypeEnum._value2member_map_[device.device_type]
                devices.setdefault(device_type, dict())[device_uuid] = device
            else:
                if not exclude_managed_by_controller:
                    device_type = DeviceTypeEnum._value2member_map_[device.device_type]
                    devices.setdefault(device_type, dict())[device_uuid] = device
                device_type = DeviceTypeEnum._value2member_map_[controller.device_type]
                devices.setdefault(device_type, dict())[controller.device_id.device_uuid.uuid] = controller
        return devices

    # ----- Service-related methods ------------------------------------------------------------------------------------

    def get_service(self, service_id : ServiceId) -> Service:
        service_key = get_service_key(service_id)
        service = self._load_grpc_object(CacheableObjectType.SERVICE, service_key)
        if service is None:
            service = get_service_by_id(self._context_client, service_id)
            if service is None: raise NotFoundException('Service', service_key)
            service : service = self._store_editable_grpc_object(
                CacheableObjectType.SERVICE, service_key, Service, service)
        return service

    def set_service(self, service : Service) -> None:
        service_key = get_service_key(service.service_id)
        self._context_client.SetService(service)
        self._store_grpc_object(CacheableObjectType.SERVICE, service_key, service)

    def delete_service(self, service_id : ServiceId) -> None:
        service_key = get_service_key(service_id)
        self._context_client.RemoveService(service_id)
        self._delete_grpc_object(CacheableObjectType.SERVICE, service_key)

    # ----- Service Handler Factory ------------------------------------------------------------------------------------

    def get_service_handlers(
        self, connection : Connection, service : Service, **service_handler_settings
    ) -> Dict[DeviceTypeEnum, Tuple['_ServiceHandler', Dict[str, Device]]]:
        connection_device_types : Dict[DeviceTypeEnum, Dict[str, Device]] = self.get_devices_from_connection(
            connection, exclude_managed_by_controller=True
        )
        service_handlers : Dict[DeviceTypeEnum, Tuple['_ServiceHandler', Dict[str, Device]]] = dict()
        for device_type, connection_devices in connection_device_types.items():
            try:
                service_handler_class = get_service_handler_class(
                    self._service_handler_factory, service, connection_devices)
                service_handler = service_handler_class(service, self, **service_handler_settings)
                service_handlers[device_type] = (service_handler, connection_devices)
            except (
                UnsatisfiedFilterException, UnsupportedFilterFieldException,
                UnsupportedFilterFieldValueException
            ):
                dict_connection_devices = {
                    cd_data.name : (cd_uuid, cd_data.name, {
                        (device_driver, DeviceDriverEnum.Name(device_driver))
                        for device_driver in cd_data.device_drivers
                    })
                    for cd_uuid,cd_data in connection_devices.items()
                }
                MSG = 'Unable to select service handler. service={:s} connection={:s} connection_devices={:s}'
                LOGGER.exception(MSG.format(
                    grpc_message_to_json_string(service), grpc_message_to_json_string(connection),
                    str(dict_connection_devices)
                ))
        return service_handlers


    # ----- QkdApp-related methods -------------------------------------------------------------------------------------

    def register_qkd_app(self, app: App) -> None:
        """
        Registers a QKD App and stores it in the cache.
        """
        qkd_app_key = get_qkd_app_key(app.app_id)
        self._qkd_app_client.RegisterApp(app)
        LOGGER.info("QKD app registered with key: %s", qkd_app_key)
        self._store_grpc_object(CacheableObjectType.QKD_APP, qkd_app_key, app)

    def update_qkd_app_status(self, app: App, new_status: QKDAppStatusEnum) -> None:
        """
        Updates the status of a QKD app and persists it to the database.
        """
        try:
            app.app_status = new_status
            LOGGER.info(f"Attempting to update app {app.app_id.app_uuid.uuid} to status {new_status}")
            self._qkd_app_client.UpdateApp(app)
            LOGGER.info(f"Successfully updated app {app.app_id.app_uuid.uuid} to status {new_status}")
        except Exception as e:
            LOGGER.error(f"Failed to update QKD app {app.app_id.app_uuid.uuid}: {str(e)}")
            raise e

    def list_qkd_apps(self, context_id: ContextId) -> List[App]:
        """
        Retrieves a list of QKD apps from the QKD App service.
        """
        try:
            apps_response = self._qkd_app_client.ListApps(context_id)
            LOGGER.info(f"ListApps retrieved: {len(apps_response.apps)} apps with status")
            
            # Ensure that the status is logged and used
            for app in apps_response.apps:
                LOGGER.info(f"App ID: {app.app_id.app_uuid.uuid}, Status: {app.app_status}")
            
            return apps_response.apps
        except Exception as e:
            LOGGER.error(f"Failed to list QKD apps: {str(e)}")
            return []

    def delete_qkd_app(self, app_id: AppId) -> None:
        """
        Deletes a QKD App by its AppId and removes it from the cache.
        """
        qkd_app_key = get_qkd_app_key(app_id)
        try:
            LOGGER.info(f"Attempting to delete QKD app with AppId: {app_id}")
            self._qkd_app_client.DeleteApp(app_id)
            LOGGER.info(f"QKD app deleted with key: {qkd_app_key}")
            self._delete_grpc_object(CacheableObjectType.QKD_APP, qkd_app_key)
        except Exception as e:
            LOGGER.error(f"Failed to delete QKD app with AppId {app_id}: {str(e)}")
            raise e

