
# Copyright © 2023-2026 Cognizant Technology Solutions Corp, www.cognizant.com.
#
# 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.
#
# END COPYRIGHT
from typing import Any
from typing import Dict

from os import environ

from asyncio import AbstractEventLoop
from asyncio import Event
from asyncio import run_coroutine_threadsafe
from concurrent.futures import ThreadPoolExecutor
from logging import getLogger
from logging import Logger

from janus import Queue

from leaf_common.config.resolver_util import ResolverUtil

from neuro_san.internals.chat.async_collating_queue import AsyncCollatingQueue
from neuro_san.internals.interfaces.reservations_storage import ReservationsStorage
from neuro_san.internals.network_providers.agent_network_storage import AgentNetworkStorage
from neuro_san.internals.network_providers.expiring_agent_network_storage import ExpiringAgentNetworkStorage
from neuro_san.internals.reservations.agent_reservation import AgentReservation
from neuro_san.internals.reservations.abstract_agent_reservationist import AbstractAgentReservationist
from neuro_san.internals.interfaces.startable import Startable


class TempNetworkStorageUpdater(Startable):
    """
    StorageUpdater implementation for temporary network being generated by neuro-san service.
    """
    # Size of thread pool for processing reservations queues.
    THREAD_POOL_SIZE: int = 32

    def __init__(self, network_storage_dict: Dict[str, AgentNetworkStorage],
                 queues: Queue[AsyncCollatingQueue]):
        """
        Constructor

        :param network_storage_dict: A dictionary of string (descripting scope) to
                    AgentNetworkStorage instance which keeps all the AgentNetwork instances
                    of a particular grouping.
        :param queues: A Queue of AsyncCollatingQueues for temp network deployment
        """
        self.logger: Logger = getLogger(self.__class__.__name__)

        self.reservationist: AbstractAgentReservationist = None
        self.queue_processing_pool: ThreadPoolExecutor = None
        self.temp_storage: ReservationsStorage = network_storage_dict.get("temp")
        if self.temp_storage is not None:
            # If we don't have temp storage, we don't got nothin'
            # Check that the temp storage instance is indeed an ExpiringReservationsStorage,
            # since we need to be able to expire reservations in it.
            if not isinstance(self.temp_storage, ExpiringAgentNetworkStorage):
                raise ValueError("Temp reservations storage must be an instance of ExpiringAgentNetworkStorage")

            # Potentially create a "base" reservations storage class
            storage_class_name: str = environ.get("AGENT_EXTERNAL_RESERVATIONS_STORAGE", "")
            external_storage: ReservationsStorage = ResolverUtil.create_instance(
                    storage_class_name,
                    "AGENT_EXTERNAL_RESERVATIONS_STORAGE env var",
                    ReservationsStorage)
            if external_storage is not None:
                self.temp_storage.set_base_storage(external_storage)
            self.reservationist = AbstractAgentReservationist({self.temp_storage})

            # We use a thread pool to process the collection of tasks
            # this TempNetworkStorageUpdater instance is responsible for:
            # 1. Checking for new queues to process coming in self.incoming "main" queue
            # 2. Processing the individual queues that come in,
            #    which involves processing the items that come in on those queues
            #    and deploying reservations to the temp storage, which is done by self.reservationist.
            self.queue_processing_pool = ThreadPoolExecutor(self.THREAD_POOL_SIZE)

            self.incoming: Queue[AsyncCollatingQueue] = queues

        else:
            # We don't have a temp storage, so we won't be doing any processing,
            # close the incoming queues and don't start any processing threads.
            self.logger.info("No temp storage configured, closing incoming queues without processing")
            queues.close()

    def start(self):
        """
        Perform start up.
        """
        if self.temp_storage is None:
            self.logger.info("TempNetworkStorageUpdater is not started: no temp storage configured.")
            return

        self.logger.info("Starting TempNetworkStorageUpdater")
        # Start any Startables
        if isinstance(self.temp_storage, Startable):
            self.temp_storage.start()
        # Start the task for checking for new queues to process,
        # this task is expected to run for the lifetime of this TempNetworkStorageUpdater instance.
        self.queue_processing_pool.submit(self.add_new_queues_to_processing)

    def add_new_queues_to_processing(self):
        """
        Checks our master queue of queues for any new additions
        and adds them to the queue processing pool
        for independent processing of each queue.
        """
        self.logger.info("Waiting for new queues to be added for processing")
        while True:
            async_collating_queue: AsyncCollatingQueue = self.incoming.sync_q.get()
            # We have a new queue to process.
            # Submit this queue processing  task to our thread pool,
            # IF we have a storage to put the results in.
            # Since each queue is processed by its own thread,
            # we avoid potential locking out of different queues received from "main" self.incoming
            if self.reservationist is not None:
                self.queue_processing_pool.submit(self.process_one_queue, async_collating_queue)
            else:
                # We don't have a reservation functionality set up,
                # so we can't do anything with this queue, but we should still close it to avoid leaks.
                self.logger.warning("No reservationist available, closing incoming queue without processing")
                async_collating_queue.close()

    def process_one_queue(self, async_collating_queue: AsyncCollatingQueue):
        """
        Processes a single AsyncCollatingQueue from the pool

        :param async_collating_queue: The AsyncCollatingQueue to process
        """
        janus_queue: Queue = async_collating_queue.get_queue()

        # See what has come over this particular queue
        while True:
            # Get an item off the queue
            queued_item: Dict[str, Any] = janus_queue.sync_q.get()
            if async_collating_queue.is_final_item(queued_item):
                # We have exhausted this queue.
                # No one needs to worry about it anymore.
                async_collating_queue.close()
                return

            self.process_one_queued_item(queued_item)

    def process_one_queued_item(self, queued_item: Dict[str, Any]):
        """
        Process a single item from one of the queues

        :param queued_item: A dictionary from a queue containing information
                as populated by AgentReservationist.deploy()
        """

        # Get the salient information from the dictionary that was queued.
        source: str = queued_item.get("source")
        deployment_dict: Dict[AgentReservation, Dict[str, Any]] = queued_item.get("deployment_dict")
        max_lifetime_in_seconds: float = queued_item.get("max_lifetime_in_seconds")

        # Do the deployment
        self.reservationist.deploy_together(deployment_dict, source, max_lifetime_in_seconds)

        # Maybe notify the deployer.
        event: Event = queued_item.get("event")
        if event is not None:
            event_loop: AbstractEventLoop = queued_item.get("event_loop")
            run_coroutine_threadsafe(self.reservationist.set_event(event), event_loop)

    def stop(self):
        """
        Perform stopping.
        """
        self.logger.info("Stopping TempNetworkStorageUpdater")
        if self.temp_storage is None:
            # We never started, so we don't need to stop anything.
            return

        # Stop any Startables
        self.temp_storage.stop()

        # Finalize the queue processing pool
        self.queue_processing_pool.shutdown(wait=False, cancel_futures=True)
