"""
In-process message bus implementations.

Provides a simple MemoryBus suitable for local development and tests.
Stub RedisBus and KafkaBus fall back to MemoryBus behavior.
"""
from __future__ import annotations

import asyncio
import uuid
from typing import Any, Dict, List


class MemoryBus:
    """A simple in-memory async message bus.

    Implements the minimal API used by Agent runtime:
      - register_agent(name, metadata)
      - publish_task(agent_name, task)
      - publish_response(task_id, response)
      - consume_tasks(agent_name, block_ms, count)
      - wait_for_response(task_id, timeout_sec)
      - get_registered_agents()
      - list_agents()
      - get_queue_depth(agent_name)
      - health_check()
    """

    def __init__(self):
        self._registry: Dict[str, Dict[str, Any]] = {}
        self._queues: Dict[str, List[Dict[str, Any]]] = {}
        self._responses: Dict[str, Any] = {}
        self._response_events: Dict[str, asyncio.Event] = {}
        self._lock = asyncio.Lock()

    async def register_agent(self, name: str, metadata: dict) -> bool:
        async with self._lock:
            self._registry[name] = {**(self._registry.get(name) or {}), **(metadata or {})}
            self._queues.setdefault(name, [])
        return True

    async def publish_task(self, agent_name: str, task: dict) -> str:
        task_id = str(uuid.uuid4())
        event = asyncio.Event()
        async with self._lock:
            self._queues.setdefault(agent_name, []).append({
                "task_id": task_id,
                "payload": task,
            })
            self._response_events[task_id] = event
        return task_id

    async def publish_response(self, task_id: str, response: dict) -> bool:
        async with self._lock:
            self._responses[task_id] = response
            event = self._response_events.get(task_id)
        if event:
            event.set()
        return True

    async def consume_tasks(self, agent_name: str, block_ms: int = 0, count: int = 1) -> list[dict]:
        # Simple non-blocking pop; if empty and block requested, sleep once
        async with self._lock:
            queue = self._queues.setdefault(agent_name, [])
            if queue:
                out, remain = queue[:count], queue[count:]
                self._queues[agent_name] = remain
                return list(out)
        # Nothing immediately available; optionally wait
        if block_ms and block_ms > 0:
            try:
                await asyncio.sleep(block_ms / 1000.0)
            except asyncio.CancelledError:
                return []
        # Try again once after waiting
        async with self._lock:
            queue = self._queues.setdefault(agent_name, [])
            if queue:
                out, remain = queue[:count], queue[count:]
                self._queues[agent_name] = remain
                return list(out)
        return []

    async def wait_for_response(self, task_id: str, timeout_sec: int = 30) -> dict | None:
        async with self._lock:
            if task_id in self._responses:
                return self._responses[task_id]
            event = self._response_events.get(task_id) or asyncio.Event()
            self._response_events[task_id] = event
        try:
            await asyncio.wait_for(event.wait(), timeout=timeout_sec)
        except asyncio.TimeoutError:
            return None
        async with self._lock:
            return self._responses.get(task_id)

    async def get_registered_agents(self) -> dict[str, dict]:
        async with self._lock:
            return dict(self._registry)

    async def list_agents(self) -> list[dict]:
        async with self._lock:
            return [{"name": name, **meta} for name, meta in self._registry.items()]

    async def get_queue_depth(self, agent_name: str) -> int:
        async with self._lock:
            return len(self._queues.get(agent_name, []))

    async def health_check(self) -> bool:
        return True


class RedisBus(MemoryBus):
    """Placeholder Redis bus using memory fallback for local runs."""
    def __init__(self, redis_url: str):  # noqa: ARG002 (unused for fallback)
        super().__init__()


class KafkaBus(MemoryBus):
    """Placeholder Kafka bus using memory fallback for local runs."""
    def __init__(self, bootstrap_servers: str):  # noqa: ARG002 (unused for fallback)
        super().__init__()
