# #region EventBusModule [C:5] [TYPE Module] [SEMANTICS event,bus,log,buffer,subscriber,flush]
# @defgroup TaskManager Module group.
# @BRIEF Task-scoped log buffering, persistence flushes, and subscriber fan-out for real-time
#        WebSocket observers — fully async with asyncio.Queue fan-out.
# @LAYER Core
# @RELATION DEPENDS_ON -> [LogEntry]
# @RELATION DEPENDS_ON -> [TaskLogPersistenceService]
# @RELATION DEPENDS_ON -> [should_log_task_level]
# @RATIONALE Extracted from TaskManager to satisfy INV_7. Migrated from threading to asyncio
#            for T043 — replaces queue.Queue + call_soon_threadsafe with asyncio.Queue
#            and async subscriber generator pattern.
# @REJECTED Keeping threading-based EventBus was rejected — call_soon_threadsafe is fragile
#            under high load and cannot be tested without a running event loop.
# @PRE TaskLogPersistenceService is initialized.
# @POST Accepted task logs are buffered, optionally persisted, and dispatched to active subscribers.
# @SIDE_EFFECT Writes task logs to persistence, mutates in-memory buffers, pushes log entries
#              into subscriber queues via async put.
# @INVARIANT Buffered logs are retried on persistence failure; every subscriber receives only
#            task-scoped events.
# @DATA_CONTRACT Input: LogEntry -> Output: persisted log + subscriber notification

import asyncio
from typing import Any

from src.core.cot_logger import seed_trace_id
from src.core.logger import belief_scope, logger, should_log_task_level
from src.core.task_manager.models import LogEntry, LogFilter, LogStats
from src.core.task_manager.persistence import TaskLogPersistenceService

QUEUE_MAXSIZE = 10000


# #region EventBus [C:5] [TYPE Class] [SEMANTICS event,bus,log,buffer,subscriber,flush]
# @defgroup TaskManager Module group.
# @BRIEF Coordinates task-scoped log buffering, persistence flushes, and subscriber fan-out
#        for real-time observers. Fully async — uses asyncio.Queue for fan-out delivery.
# @RELATION DEPENDS_ON -> [LogEntry]
# @RELATION DEPENDS_ON -> [TaskLogPersistenceService]
# @RELATION DEPENDS_ON -> [should_log_task_level]
# @PRE TaskLogPersistenceService is initialized.
# @POST Accepted task logs are buffered, optionally persisted, and dispatched to active subscribers.
# @INVARIANT Buffered logs are retried on persistence failure.
# @RATIONALE Fully async event bus with asyncio.Queue(maxsize=10000) for each subscriber.
#            Broadcast methods use await queue.put() instead of call_soon_threadsafe.
# @REJECTED Keeping threading-based EventBus with call_soon_threadsafe — fragile under load,
#            hard to test, and incompatible with fully async task execution.
class EventBus:
    """Task-scoped log buffering, persistence flushes, and subscriber fan-out — fully async."""

    LOG_FLUSH_INTERVAL = 2.0

    # #region __init__ [TYPE Function]
    # @BRIEF Initialize log buffer, subscriber registries, and async flusher.
    # @POST EventBus is ready to accept logs. Subscriber dicts are empty.
    def __init__(self, log_persistence_service: TaskLogPersistenceService):
        self.log_persistence_service = log_persistence_service
        self.subscribers: dict[str, list[asyncio.Queue]] = {}
        self._status_subscribers: dict[str, list[asyncio.Queue]] = {}
        self._task_event_subscribers: list[asyncio.Queue] = []
        self._maintenance_subscribers: list[asyncio.Queue] = []
        self._log_buffer: dict[str, list[LogEntry]] = {}
        self._flusher_stop_event = asyncio.Event()
        self._flusher_task: asyncio.Task | None = None
    # #endregion __init__

    # #region start [TYPE Function]
# @ingroup TaskManager
    # @BRIEF Start the background async flusher task.
    # @POST Flusher task is created and scheduled in the event loop.
    def start(self) -> None:
        if self._flusher_task is None or self._flusher_task.done():
            self._flusher_task = asyncio.create_task(self.async_flusher_loop())
    # #endregion start

    # #region stop [TYPE Function]
# @ingroup TaskManager
    # @BRIEF Signal the flusher to stop and wait for it.
    # @POST Flusher task is cancelled and awaited.
    async def stop(self) -> None:
        self._flusher_stop_event.set()
        if self._flusher_task and not self._flusher_task.done():
            self._flusher_task.cancel()
            from contextlib import suppress
            with suppress(asyncio.CancelledError, Exception):
                await self._flusher_task
            self._flusher_task = None
    # #endregion stop

    # #region async_flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,async]
# @ingroup TaskManager
    # @BRIEF Background async task that periodically flushes log buffer to database.
    async def async_flusher_loop(self):
        seed_trace_id()
        try:
            while not self._flusher_stop_event.is_set():
                await self._flush_logs()
                try:
                    await asyncio.wait_for(
                        self._flusher_stop_event.wait(),
                        timeout=self.LOG_FLUSH_INTERVAL,
                    )
                    break  # Event was set
                except TimeoutError:
                    continue  # Timeout is normal — flush again
        except asyncio.CancelledError:
            pass
    # #endregion async_flusher_loop

    # #region _flush_logs [C:3] [TYPE Function] [SEMANTICS flush,batch,persistence]
    # @BRIEF Flush all buffered logs to the database.
    # @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.add_logs]
    async def _flush_logs(self):
        seed_trace_id()
        task_ids = list(self._log_buffer.keys())
        for task_id in task_ids:
            logs = self._log_buffer.pop(task_id, [])
            if logs:
                try:
                    self.log_persistence_service.add_logs(task_id, logs)
                    logger.debug(f"Flushed {len(logs)} logs for task {task_id}")
                except Exception as e:
                    logger.error(f"Failed to flush logs for task {task_id}: {e}")
                    if task_id not in self._log_buffer:
                        self._log_buffer[task_id] = []
                    self._log_buffer[task_id].extend(logs)
    # #endregion _flush_logs

    # #region flush_task_logs [C:3] [TYPE Function] [SEMANTICS flush,single,persistence]
# @ingroup TaskManager
    # @BRIEF Flush logs for a specific task immediately.
    # @PRE task_id exists.
    # @POST Task's buffered logs are written to database.
    # @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.add_logs]
    async def flush_task_logs(self, task_id: str):
        with belief_scope("EventBus.flush_task_logs"):
            logs = self._log_buffer.pop(task_id, [])
            if logs:
                try:
                    self.log_persistence_service.add_logs(task_id, logs)
                except Exception as e:
                    logger.error(f"Failed to flush logs for task {task_id}: {e}")
    # #endregion flush_task_logs

    # #region add_log [C:3] [TYPE Function] [SEMANTICS log,buffer,notify,subscriber]
# @ingroup TaskManager
    # @BRIEF Adds a log entry to a task buffer and notifies subscribers via async put.
    # @PRE Task exists.
    # @POST Log added to buffer and pushed to subscriber queues (if level passes filter).
    # @RELATION CALLS -> [should_log_task_level]
    async def add_log(
        self,
        task_id: str,
        level: str,
        message: str,
        source: str = "system",
        metadata: dict[str, Any] | None = None,
        context: dict[str, Any] | None = None,
        task_logs_list: list[LogEntry] | None = None,
    ):
        # Filter logs based on task_log_level configuration
        if not should_log_task_level(level):
            return
        log_entry = LogEntry(
            level=level,
            message=message,
            source=source,
            metadata=metadata,
            context=context,
        )
        # Add to in-memory logs (backward compatibility)
        if task_logs_list is not None:
            task_logs_list.append(log_entry)
        # Add to buffer for batch persistence
        if task_id not in self._log_buffer:
            self._log_buffer[task_id] = []
        self._log_buffer[task_id].append(log_entry)
        # Notify subscribers via async put (fan-out)
        if task_id in self.subscribers:
            for queue in self.subscribers[task_id]:
                try:
                    await queue.put(log_entry)
                except asyncio.QueueFull:
                    logger.warning(
                        f"Subscriber queue full for task {task_id}, dropping log entry"
                    )
    # #endregion add_log

    # ── Log Subscribers ──

    # #region subscribe_logs [C:2] [TYPE Function] [SEMANTICS subscribe,queue,realtime]
# @ingroup TaskManager
    # @BRIEF Subscribes to real-time logs for a task.
    # @POST Returns an asyncio.Queue (maxsize=10000) for log entries.
    async def subscribe_logs(self, task_id: str) -> asyncio.Queue:
        queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
        if task_id not in self.subscribers:
            self.subscribers[task_id] = []
        self.subscribers[task_id].append(queue)
        return queue
    # #endregion subscribe_logs

    # #region unsubscribe_logs [C:2] [TYPE Function] [SEMANTICS unsubscribe,queue,remove]
# @ingroup TaskManager
    # @BRIEF Unsubscribes from real-time logs for a task.
    # @POST Queue removed from subscribers.
    def unsubscribe_logs(self, task_id: str, queue: asyncio.Queue):
        if task_id in self.subscribers:
            if queue in self.subscribers[task_id]:
                self.subscribers[task_id].remove(queue)
            if not self.subscribers[task_id]:
                del self.subscribers[task_id]
    # #endregion unsubscribe_logs

    # ── Task Status Subscribers ──

    # #region subscribe_status [C:2] [TYPE Function] [SEMANTICS subscribe,status,realtime]
# @ingroup TaskManager
    # @BRIEF Subscribes to real-time status updates for a specific task.
    # @POST Returns an asyncio.Queue (maxsize=10000) for Task status dicts.
    async def subscribe_status(self, task_id: str) -> asyncio.Queue:
        queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
        if task_id not in self._status_subscribers:
            self._status_subscribers[task_id] = []
        self._status_subscribers[task_id].append(queue)
        return queue
    # #endregion subscribe_status

    # #region unsubscribe_status [C:2] [TYPE Function] [SEMANTICS unsubscribe,status,remove]
# @ingroup TaskManager
    # @BRIEF Unsubscribes from status updates for a task.
    # @POST Queue removed from status subscribers.
    def unsubscribe_status(self, task_id: str, queue: asyncio.Queue):
        if task_id in self._status_subscribers:
            if queue in self._status_subscribers[task_id]:
                self._status_subscribers[task_id].remove(queue)
            if not self._status_subscribers[task_id]:
                del self._status_subscribers[task_id]
    # #endregion unsubscribe_status

    # #region broadcast_status [C:2] [TYPE Function] [SEMANTICS broadcast,status,notify]
# @ingroup TaskManager
    # @BRIEF Broadcast a task status update to all per-task and global subscribers.
    # @POST Status event pushed to all matching queues via async put.
    async def broadcast_status(self, task_id: str, task_dict: dict):
        event = {"type": "task_status", "task_id": task_id, "task": task_dict}
        # Per-task subscribers
        if task_id in self._status_subscribers:
            for queue in self._status_subscribers[task_id]:
                try:
                    await queue.put(event)
                except asyncio.QueueFull:
                    logger.warning(
                        f"Status subscriber queue full for task {task_id}, dropping event"
                    )
        # Global task event subscribers
        for queue in self._task_event_subscribers:
            try:
                await queue.put(event)
            except asyncio.QueueFull:
                logger.warning(
                    "Global task event subscriber queue full, dropping event"
                )
    # #endregion broadcast_status

    # ── Maintenance Event Subscribers ──

    # #region subscribe_maintenance_events [C:2] [TYPE Function] [SEMANTICS subscribe,maintenance,events]
# @ingroup TaskManager
    # @BRIEF Subscribes to maintenance event updates (created/ended/banner changes).
    # @POST Returns an asyncio.Queue (maxsize=10000) for maintenance event dicts.
    async def subscribe_maintenance_events(self) -> asyncio.Queue:
        queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
        self._maintenance_subscribers.append(queue)
        return queue
    # #endregion subscribe_maintenance_events

    # #region unsubscribe_maintenance_events [C:2] [TYPE Function] [SEMANTICS unsubscribe,maintenance,events]
# @ingroup TaskManager
    # @BRIEF Unsubscribes from maintenance events.
    # @POST Queue removed from maintenance subscribers.
    def unsubscribe_maintenance_events(self, queue: asyncio.Queue):
        if queue in self._maintenance_subscribers:
            self._maintenance_subscribers.remove(queue)
    # #endregion unsubscribe_maintenance_events

    # #region broadcast_maintenance_event [C:2] [TYPE Function] [SEMANTICS broadcast,maintenance,notify]
# @ingroup TaskManager
    # @BRIEF Broadcast a maintenance event to all subscribers.
    # @POST Event pushed to all maintenance subscriber queues via async put.
    async def broadcast_maintenance_event(self, event: dict):
        for queue in self._maintenance_subscribers:
            try:
                await queue.put(event)
            except asyncio.QueueFull:
                logger.warning(
                    "Maintenance subscriber queue full, dropping event"
                )
    # #endregion broadcast_maintenance_event

    # ── Global Task Event Subscribers ──

    # #region subscribe_task_events [C:2] [TYPE Function] [SEMANTICS subscribe,global,events]
# @ingroup TaskManager
    # @BRIEF Subscribes to ALL task events (status changes) globally.
    # @POST Returns an asyncio.Queue (maxsize=10000) for task event dicts.
    async def subscribe_task_events(self) -> asyncio.Queue:
        queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
        self._task_event_subscribers.append(queue)
        return queue
    # #endregion subscribe_task_events

    # #region unsubscribe_task_events [C:2] [TYPE Function] [SEMANTICS unsubscribe,global,events]
# @ingroup TaskManager
    # @BRIEF Unsubscribes from global task events.
    # @POST Queue removed from global subscribers.
    def unsubscribe_task_events(self, queue: asyncio.Queue):
        if queue in self._task_event_subscribers:
            self._task_event_subscribers.remove(queue)
    # #endregion unsubscribe_task_events

    # ── Log retrieval methods (sync, delegates to persistence) ──

    # #region get_task_logs [C:3] [TYPE Function] [SEMANTICS logs,read,persistence,backfill]
# @ingroup TaskManager
    # @BRIEF Retrieves logs for a specific task (from memory for running, persistence for completed).
    # @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.get_logs]
    def get_task_logs(
        self, task_id: str, log_filter: LogFilter | None = None, task_status=None, task_logs: list | None = None
    ) -> list[LogEntry]:
        from src.core.task_manager.models import TaskStatus
        # For completed tasks, fetch from persistence
        if task_status in [TaskStatus.SUCCESS, TaskStatus.FAILED]:
            if log_filter is None:
                log_filter = LogFilter()
            task_logs_persisted = self.log_persistence_service.get_logs(task_id, log_filter)
            return [
                LogEntry(
                    timestamp=log.timestamp,
                    level=log.level,
                    message=log.message,
                    source=log.source,
                    metadata=log.metadata,
                )
                for log in task_logs_persisted
            ]
        # For running/pending tasks, return from memory
        return task_logs if task_logs else []
    # #endregion get_task_logs

    # #region get_task_log_stats [C:2] [TYPE Function] [SEMANTICS log,stats,aggregate]
# @ingroup TaskManager
    # @BRIEF Get statistics about logs for a task.
    # @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.get_log_stats]
    def get_task_log_stats(self, task_id: str) -> LogStats:
        return self.log_persistence_service.get_log_stats(task_id)
    # #endregion get_task_log_stats

    # #region get_task_log_sources [C:2] [TYPE Function] [SEMANTICS log,sources,unique]
# @ingroup TaskManager
    # @BRIEF Get unique sources for a task's logs.
    # @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.get_sources]
    def get_task_log_sources(self, task_id: str) -> list[str]:
        return self.log_persistence_service.get_sources(task_id)
    # #endregion get_task_log_sources

    # #region delete_logs_for_tasks [C:2] [TYPE Function] [SEMANTICS log,delete,cleanup]
# @ingroup TaskManager
    # @BRIEF Delete logs for specified tasks.
    # @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.delete_logs_for_tasks]
    def delete_logs_for_tasks(self, task_ids: list[str]) -> None:
        if task_ids:
            self.log_persistence_service.delete_logs_for_tasks(task_ids)
    # #endregion delete_logs_for_tasks
# #endregion EventBus
# #endregion EventBusModule
