
# 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 asyncio import sleep as asyncio_sleep
from asyncio import CancelledError
from logging import getLogger
from logging import Logger
import queue as queue_mod

from janus import Queue
from janus import SyncQueueShutDown

from leaf_common.config.resolver_util import ResolverUtil
from leaf_common.asyncio.asyncio_executor import AsyncioExecutor

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

    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.executor: AsyncioExecutor = 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 an AsyncioExecutor and its event loop 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.executor = AsyncioExecutor()

            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()

        if self.executor is not None:
            self.executor.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.executor.create_task(self.add_new_queues_to_processing(), "incoming queue")

    async def add_new_queues_to_processing(self):
        """
        Checks our master queue of queues for any new additions
        and adds processing tasks to executor
        for independent processing of each queue.
        """
        idle_sleep_seconds: float = 2.0

        self.logger.info("Waiting for new queues to be added for processing")
        while True:
            try:
                async_collating_queue: AsyncCollatingQueue = self.incoming.sync_q.get_nowait()
            except queue_mod.Empty:
                # No item to process right now, sleep for a bit and check again.
                # But first check, if the queue is closed already:
                if self.incoming.sync_q.closed:
                    self.logger.info("Incoming queue closed")
                    return
                try:
                    await asyncio_sleep(idle_sleep_seconds)
                except CancelledError:
                    self.logger.info("Incoming queue task cancelled during sleep")
                    raise
                # Continue polling for new queues to process.
                continue
            except SyncQueueShutDown:
                # Explicit shutdown signal from Janus queue machinery
                self.logger.info("SHUTDOWN signal for incoming queue")
                return

            if self.reservationist is not None:
                _ = self.executor.create_task(self.process_one_queue(async_collating_queue), "queue_processor")
            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()

    async def process_one_queue(self, async_collating_queue: AsyncCollatingQueue):
        """
        Processes a single AsyncCollatingQueue

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

        # See what has come over this particular queue
        while True:
            # Get an item off the queue
            try:
                queued_item: Dict[str, Any] = janus_queue.sync_q.get_nowait()
            except queue_mod.Empty:
                # No item to process right now, sleep for a bit and check again.
                # But first check, if the queue is closed already:
                if janus_queue.sync_q.closed:
                    return
                try:
                    await asyncio_sleep(idle_sleep_seconds)
                except CancelledError:
                    self.logger.info("Reservation queue task cancelled during sleep")
                    raise
                continue
            except SyncQueueShutDown:
                # Explicit shutdown signal from Janus queue machinery
                self.logger.info("SHUTDOWN signal for reservation queue")
                return
            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

            # Now: we have an item to process from this queue.
            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 not None:
            self.temp_storage.stop()

        # Finalize the executor:
        if self.executor is not None:
            self.executor.shutdown()
