"""Background module manager with single instance."""

import asyncio
import uuid
from collections.abc import AsyncGenerator, AsyncIterator
from contextlib import asynccontextmanager
from typing import Any, Generic

import grpc

from digitalkin.logger import logger
from digitalkin.models import ModuleStatus
from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
from digitalkin.modules._base_module import BaseModule
from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
from digitalkin.modules.job_manager.job_manager_models import StreamCodeModel
from digitalkin.services.services_models import ServicesMode


class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
    """Manages a single instance of a module job.

    This class ensures that only one instance of a module job is active at a time.
    It provides functionality to create, stop, and monitor module jobs, as well as
    to handle their output data.
    """

    modules: dict[str, BaseModule]
    queue: dict[str, asyncio.Queue]

    def __init__(
        self,
        module_class: type[BaseModule],
        services_mode: ServicesMode,
    ) -> None:
        """Initialize the job manager.

        Args:
            module_class: The class of the module to be managed.
            services_mode: The mode of operation for the services (e.g., ASYNC or SYNC).
        """
        super().__init__(module_class, services_mode)

        self._lock = asyncio.Lock()
        self.modules: dict[str, BaseModule] = {}
        self.queues: dict[str, asyncio.Queue] = {}

    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
        """Generate a stream consumer for a module's output data.

        This method creates an asynchronous generator that streams output data
        from a specific module job. If the module does not exist, it generates
        an error message.

        Args:
            job_id: The unique identifier of the job.

        Returns:
            SetupModelT: the SetupModelT object fully processed.
        """
        module = self.modules.get(job_id, None)
        logger.debug("Module %s found: %s", job_id, module)

        try:
            return await self.queues[job_id].get()
        finally:
            logger.info(f"{job_id=}: {self.queues[job_id].empty()}")
            del self.queues[job_id]

    async def create_config_setup_instance_job(
        self,
        config_setup_data: SetupModelT,
        mission_id: str,
        setup_id: str,
        setup_version_id: str,
    ) -> str:
        """Create and start a new module setup configuration job.

        This method initializes a new module job, assigns it a unique job ID,
        and starts the config setup it in the background.

        Args:
            config_setup_data: The input data required to start the job.
            setup_data: The setup configuration for the module.
            mission_id: The mission ID associated with the job.
            setup_id: The setup ID associated with the module.
            setup_version_id: The setup ID.

        Returns:
            str: The unique identifier (job ID) of the created job.

        Raises:
            Exception: If the module fails to start.
        """
        job_id = str(uuid.uuid4())
        # TODO: Ensure the job_id is unique.
        module = self.module_class(job_id, mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id)
        self.modules[job_id] = module
        self.queues[job_id] = asyncio.Queue()

        try:
            await module.start_config_setup(
                config_setup_data,
                await self.job_specific_callback(self.add_to_queue, job_id),
            )
            logger.debug("Module %s (%s) started successfully", job_id, module.name)
        except Exception:
            # Remove the module from the manager in case of an error.
            del self.modules[job_id]
            logger.exception("Failed to start module %s: %s", job_id)
            raise
        else:
            return job_id

    async def add_to_queue(self, job_id: str, output_data: OutputModelT) -> None:  # type: ignore
        """Add output data to the queue for a specific job.

        This method is used as a callback to handle output data generated by a module job.

        Args:
            job_id: The unique identifier of the job.
            output_data: The output data produced by the job.
        """
        await self.queues[job_id].put(output_data.model_dump())

    @asynccontextmanager  # type: ignore
    async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # type: ignore
        """Generate a stream consumer for a module's output data.

        This method creates an asynchronous generator that streams output data
        from a specific module job. If the module does not exist, it generates
        an error message.

        Args:
            job_id: The unique identifier of the job.

        Yields:
            AsyncGenerator: A stream of output data or error messages.
        """
        module = self.modules.get(job_id, None)

        logger.debug("Module %s found: %s", job_id, module)

        async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
            """Stream output data from the module.

            Yields:
                dict: Output data generated by the module.
            """
            if module is None:
                yield {
                    "error": {
                        "error_message": f"Module {job_id} not found",
                        "code": grpc.StatusCode.NOT_FOUND,
                    }
                }
                return

            try:
                while module.status == ModuleStatus.RUNNING or (
                    not self.queues[job_id].empty()
                    and module.status
                    in {
                        ModuleStatus.STOPPED,
                        ModuleStatus.STOPPING,
                    }
                ):
                    logger.info(f"{job_id=}: {module.status=}")
                    yield await self.queues[job_id].get()

            finally:
                del self.queues[job_id]

        yield _stream()

    async def create_module_instance_job(
        self,
        input_data: InputModelT,
        setup_data: SetupModelT,
        mission_id: str,
        setup_id: str,
        setup_version_id: str,
    ) -> str:
        """Create and start a new module job.

        This method initializes a new module job, assigns it a unique job ID,
        and starts it in the background.

        Args:
            input_data: The input data required to start the job.
            setup_data: The setup configuration for the module.
            mission_id: The mission ID associated with the job.
            setup_id: The setup ID associated with the module.
            setup_version_id: The setup Version ID associated with the module.

        Returns:
            str: The unique identifier (job ID) of the created job.

        Raises:
            Exception: If the module fails to start.
        """
        job_id = str(uuid.uuid4())
        # TODO: Ensure the job_id is unique.
        module = self.module_class(
            job_id,
            mission_id=mission_id,
            setup_id=setup_id,
            setup_version_id=setup_version_id,
        )
        self.modules[job_id] = module
        self.queues[job_id] = asyncio.Queue()
        callback = await self.job_specific_callback(self.add_to_queue, job_id)

        try:
            await module.start(
                input_data,
                setup_data,
                callback,
                done_callback=lambda _: asyncio.create_task(callback(StreamCodeModel(code="__END_OF_STREAM__"))),
            )
            logger.debug("Module %s (%s) started successfully", job_id, module.name)
        except Exception:
            # Remove the module from the manager in case of an error.
            del self.modules[job_id]
            logger.exception("Failed to start module %s: %s", job_id)
            raise
        else:
            return job_id

    async def stop_module(self, job_id: str) -> bool:
        """Stop a running module job.

        Args:
            job_id: The unique identifier of the job to stop.

        Returns:
            bool: True if the module was successfully stopped, False if it does not exist.

        Raises:
            Exception: If an error occurs while stopping the module.
        """
        async with self._lock:
            module = self.modules.get(job_id)
            if not module:
                logger.warning(f"Module {job_id} not found")
                return False
            try:
                await module.stop()
                # should maybe be added in finally
                del self.queues[job_id]
                del self.modules[job_id]
                logger.debug(f"Module {job_id} ({module.name}) stopped successfully")
            except Exception as e:
                logger.error(f"Error while stopping module {job_id}: {e}")
                raise
            else:
                return True

    async def get_module_status(self, job_id: str) -> ModuleStatus | None:
        """Retrieve the status of a module job.

        Args:
            job_id: The unique identifier of the job.

        Returns:
            ModuleStatus | None: The status of the module, or None if it does not exist.
        """
        module = self.modules.get(job_id)
        return module.status if module else None

    async def stop_all_modules(self) -> None:
        """Stop all currently running module jobs.

        This method ensures that all active jobs are gracefully terminated.
        """
        async with self._lock:
            stop_tasks = [self.stop_module(job_id) for job_id in list(self.modules.keys())]
            if stop_tasks:
                await asyncio.gather(*stop_tasks, return_exceptions=True)

    async def list_modules(self) -> dict[str, dict[str, Any]]:
        """List all modules along with their statuses.

        Returns:
            dict[str, dict[str, Any]]: A dictionary containing information about all modules and their statuses.
        """
        return {
            job_id: {
                "name": module.name,
                "status": module.status,
                "class": module.__class__.__name__,
            }
            for job_id, module in self.modules.items()
        }
