from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from lfx.log.logger import logger

from langflow.events.event_manager import EventManager
from langflow.services.base import Service

if TYPE_CHECKING:
    from uuid import UUID


class JobQueueNotFoundError(Exception):
    """Exception raised when a job queue is not found."""

    def __init__(self, job_id: str) -> None:
        self.job_id = job_id
        super().__init__(f"Job queue not found for job_id: {job_id}")


class JobQueueService(Service):
    """Asynchronous service for managing job-specific queues and their associated tasks.

    This service allows clients to:
      - Create dedicated asyncio queues for individual jobs.
      - Associate each queue with an EventManager, enabling event-driven handling.
      - Launch and manage asynchronous tasks that process these job queues.
      - Safely clean up resources by cancelling active tasks and emptying queues.
      - Automatically perform periodic cleanup of inactive or completed job queues.

    The cleanup process follows a two-phase approach:
      1. When a task is cancelled or fails, it is marked for cleanup by setting a timestamp
      2. The actual cleanup only occurs after CLEANUP_GRACE_PERIOD seconds have elapsed
         since the task was marked

    Attributes:
        name (str): Unique identifier for the service.
        _queues (dict[str, tuple[asyncio.Queue, EventManager, asyncio.Task | None, float | None]]):
            Dictionary mapping job IDs to a tuple containing:
              * The job's asyncio.Queue instance.
              * The associated EventManager instance.
              * The asyncio.Task processing the job (if any).
              * The cleanup timestamp (if any).
        _cleanup_task (asyncio.Task | None): Background task for periodic cleanup.
        _closed (bool): Flag indicating whether the service is currently active.
        CLEANUP_GRACE_PERIOD (int): Number of seconds to wait after a task is marked for cleanup
            before actually removing it. This grace period allows for:
              * Pending operations to complete
              * Related systems to finish their work
              * Inspection or recovery if needed
            Default is 300 seconds (5 minutes).

    Example:
        service = JobQueueService()
        await service.start()
        queue, event_manager = service.create_queue("job123")
        service.start_job("job123", some_async_coroutine())
        # Retrieve and use the queue data as needed
        data = service.get_queue_data("job123")
        await service.cleanup_job("job123")
        await service.stop()
    """

    name = "job_queue_service"

    def __init__(self) -> None:
        """Initialize the JobQueueService.

        Sets up the internal registry for job queues, initializes the cleanup task, and sets the service state
        to active.
        """
        self._queues: dict[str, tuple[asyncio.Queue, EventManager, asyncio.Task | None, float | None]] = {}
        self._job_owners: dict[str, UUID] = {}
        self._cleanup_task: asyncio.Task | None = None
        self._closed = False
        self.ready = False
        self.CLEANUP_GRACE_PERIOD = 300  # 5 minutes before cleaning up marked tasks

    def is_started(self) -> bool:
        """Check if the JobQueueService has started.

        Returns:
            bool: True if the service has started, False otherwise.
        """
        return self._cleanup_task is not None

    def set_ready(self) -> None:
        if not self.is_started():
            self.start()
        super().set_ready()

    def start(self) -> None:
        """Start the JobQueueService and begin the periodic cleanup routine.

        This method marks the service as active and launches a background task that
        periodically checks and cleans up job queues whose tasks have been completed or cancelled.
        """
        self._closed = False
        self._cleanup_task = asyncio.create_task(self._periodic_cleanup())
        logger.debug("JobQueueService started: periodic cleanup task initiated.")

    async def stop(self) -> None:
        """Gracefully stop the JobQueueService by terminating background operations and cleaning up all resources.

        This coroutine performs the following steps:
            1. Marks the service as closed, preventing further job queue creation.
            2. Cancels the background periodic cleanup task and awaits its termination.
            3. Iterates over all registered job queues to clean up their resources—cancelling active tasks and
            clearing queued items.
        """
        self._closed = True
        if self._cleanup_task:
            self._cleanup_task.cancel()
            await asyncio.wait([self._cleanup_task])
            if not self._cleanup_task.cancelled():
                exc = self._cleanup_task.exception()
                if exc is not None:
                    raise exc

        # Clean up each registered job queue.
        for job_id in list(self._queues.keys()):
            await self.cleanup_job(job_id)
        await logger.adebug("JobQueueService stopped: all job queues have been cleaned up.")

    async def teardown(self) -> None:
        await self.stop()

    def create_queue(self, job_id: str) -> tuple[asyncio.Queue, EventManager]:
        """Create and register a new queue along with its corresponding event manager for a job.

        Args:
            job_id (str): Unique identifier for the job.

        Returns:
            tuple[asyncio.Queue, EventManager]: A tuple containing:
                - The asyncio.Queue instance for handling the job's tasks or messages.
                - The EventManager instance for event handling tied to the queue.
        """
        if self._closed:
            msg = "Queue service is closed"
            raise RuntimeError(msg)

        existing_queue = self._queues.get(job_id)
        if existing_queue:
            msg = f"Queue for job_id {job_id} already exists"
            raise ValueError(msg)

        main_queue: asyncio.Queue = asyncio.Queue()
        event_manager: EventManager = self._create_default_event_manager(main_queue)

        # Register the queue without an active task.
        self._queues[job_id] = (main_queue, event_manager, None, None)
        logger.debug(f"Queue and event manager successfully created for job_id {job_id}")
        return main_queue, event_manager

    def start_job(self, job_id: str, task_coro) -> None:
        """Start an asynchronous task for a given job, replacing any existing active task.

        The method performs the following:
          - Verifies the presence of a registered queue for the job.
          - Cancels any currently running task associated with the job.
          - Launches a new asynchronous task using the provided coroutine.
          - Updates the internal registry with the new task.

        Args:
            job_id (str): Unique identifier for the job.
            task_coro: A coroutine representing the job's asynchronous task.
        """
        if job_id not in self._queues:
            msg = f"No queue found for job_id {job_id}"
            logger.error(msg)
            raise ValueError(msg)

        if self._closed:
            msg = "Queue service is closed"
            logger.error(msg)
            raise RuntimeError(msg)

        main_queue, event_manager, existing_task, _ = self._queues[job_id]
        if existing_task and not existing_task.done():
            logger.debug(f"Existing task for job_id {job_id} detected; cancelling it.")
            existing_task.cancel()

        # Initiate the new asynchronous task.
        task = asyncio.create_task(task_coro)
        self._queues[job_id] = (main_queue, event_manager, task, None)
        logger.debug(f"New task started for job_id {job_id}")

    def get_queue_data(self, job_id: str) -> tuple[asyncio.Queue, EventManager, asyncio.Task | None, float | None]:
        """Retrieve the complete data structure associated with a job's queue.

        Args:
            job_id (str): Unique identifier for the job.

        Returns:
            tuple[asyncio.Queue, EventManager, asyncio.Task | None, float | None]:
                A tuple containing the job's main queue, its linked event manager, the associated task (if any),
                and the cleanup timestamp (if any).

        Raises:
            JobQueueNotFoundError: If the job_id is not found.
            RuntimeError: If the service is closed.
        """
        if self._closed:
            msg = f"Queue service is closed for job_id: {job_id}"
            raise RuntimeError(msg)

        try:
            return self._queues[job_id]
        except KeyError as exc:
            raise JobQueueNotFoundError(job_id) from exc

    def register_job_owner(self, job_id: str, user_id: UUID) -> None:
        """Register the authenticated user who initiated a build job."""
        self._job_owners[job_id] = user_id

    def get_job_owner(self, job_id: str) -> UUID | None:
        """Return the user ID that owns a job, or None if not tracked."""
        return self._job_owners.get(job_id)

    async def cleanup_job(self, job_id: str) -> None:
        """Clean up and release resources for a specific job.

        The cleanup process includes:
          1. Verifying if the job's queue is registered.
          2. Cancelling the running task (if active) and awaiting its termination.
          3. Clearing all items from the job's queue.
          4. Removing the job's entry from the internal registry.

        Args:
            job_id (str): Unique identifier for the job to be cleaned up.
        """
        if job_id not in self._queues:
            await logger.adebug(f"No queue found for job_id {job_id} during cleanup.")
            return

        await logger.adebug(f"Commencing cleanup for job_id {job_id}")
        main_queue, _event_manager, task, _ = self._queues[job_id]

        # Cancel the associated task if it is still running.
        if task and not task.done():
            await logger.adebug(f"Cancelling active task for job_id {job_id}")
            task.cancel()
            try:
                await task
            except asyncio.CancelledError as exc:
                # Check if this was a user-initiated cancellation (user called task.cancel())
                if task.cancelled():
                    # User-initiated cancellation so we explicitly called task.cancel() above
                    await logger.adebug(f"Task for job_id {job_id} was successfully cancelled.")
                    # Re-raise with user cancellation message code
                    exc.args = ("LANGFLOW_USER_CANCELLED",)
                    raise
                # System-initiated cancellation for other reasons
                await logger.adebug(f"Task for job_id {job_id} was cancelled by system.")
                exc.args = ("LANGFLOW_SYSTEM_CANCELLED",)
                raise
            except Exception as exc:
                await logger.aerror(f"Error in task for job_id {job_id} during cancellation: {exc}")
                raise
        await logger.adebug(f"Task cancellation complete for job_id {job_id}")

        # Clear the queue since we just cancelled the task or it has completed
        items_cleared = 0
        while not main_queue.empty():
            try:
                main_queue.get_nowait()
                items_cleared += 1
            except asyncio.QueueEmpty:
                break

        await logger.adebug(f"Removed {items_cleared} items from queue for job_id {job_id}")
        # Remove the job entry from the registry
        self._queues.pop(job_id, None)
        self._job_owners.pop(job_id, None)
        await logger.adebug(f"Cleanup successful for job_id {job_id}: resources have been released.")

    async def _periodic_cleanup(self) -> None:
        """Execute a periodic task that cleans up completed or cancelled job queues.

        This internal coroutine continuously:
          - Sleeps for a fixed interval (60 seconds).
          - Initiates the cleanup of job queues by calling _cleanup_old_queues.
          - Monitors and logs any exceptions during the cleanup cycle.

        The loop terminates when the service is marked as closed.
        """
        while not self._closed:
            try:
                await asyncio.sleep(60)  # Sleep for 60 seconds before next cleanup attempt.
                await self._cleanup_old_queues()
            except asyncio.CancelledError:
                await logger.adebug("Periodic cleanup task received cancellation signal.")
                raise
            except Exception as exc:  # noqa: BLE001
                await logger.aerror(f"Exception encountered during periodic cleanup: {exc}")

    async def _cleanup_old_queues(self) -> None:
        """Scan all registered job queues and clean up those with completed, failed or orphaned tasks."""
        current_time = asyncio.get_running_loop().time()

        for job_id in list(self._queues.keys()):
            _, _, task, cleanup_time = self._queues[job_id]

            should_cleanup = False
            cleanup_reason = ""

            # Case 1: Orphaned queue (created but task never started)
            if task is None:
                should_cleanup = True
                cleanup_reason = "Orphaned queue (no task associated)"
            # Case 2: Task has finished (Success, Failure, or Cancellation)
            elif task.done():
                should_cleanup = True
                if task.cancelled():
                    cleanup_reason = "Task cancelled"
                elif task.exception() is not None:
                    # Don't try to log the exception yet as it might be handled elsewhere;
                    # the grace period allows other systems to inspect it if needed.
                    cleanup_reason = "Task failed with exception"
                else:
                    cleanup_reason = "Task completed successfully"

            if should_cleanup:
                if cleanup_time is None:
                    # Mark for cleanup by setting the timestamp
                    self._queues[job_id] = (
                        self._queues[job_id][0],
                        self._queues[job_id][1],
                        self._queues[job_id][2],
                        current_time,
                    )
                    await logger.adebug(f"Job queue for job_id {job_id} marked for cleanup - {cleanup_reason}")
                elif current_time - cleanup_time >= self.CLEANUP_GRACE_PERIOD:
                    # Enough time has passed, perform the actual cleanup
                    await logger.adebug(f"Cleaning up job_id {job_id} after grace period due to: {cleanup_reason}")
                    await self.cleanup_job(job_id)

    def _create_default_event_manager(self, queue: asyncio.Queue) -> EventManager:
        """Creates the default event manager with predefined events.

        Args:
            queue (asyncio.Queue): The queue to be associated with the event manager.

        Returns:
            EventManager: The configured EventManager instance.
        """
        manager = EventManager(queue)
        # Registering predefined events
        event_names_types = [
            ("on_token", "token"),
            ("on_vertices_sorted", "vertices_sorted"),
            ("on_error", "error"),
            ("on_end", "end"),
            ("on_message", "add_message"),
            ("on_remove_message", "remove_message"),
            ("on_end_vertex", "end_vertex"),
            ("on_build_start", "build_start"),
            ("on_build_end", "build_end"),
            ("on_log", "log"),
        ]
        for name, event_type in event_names_types:
            manager.register_event(name, event_type)
        return manager
