"""
Message Bus for Inter-Agent Communication

Provides:
- Pub/Sub messaging via Redis/Valkey
- Request-Response pattern with timeouts
- Priority-based message handling
- Message persistence for audit
"""

from typing import Callable, Dict, Any, List, Optional, Awaitable
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
import asyncio
import uuid
import json
import logging

logger = logging.getLogger(__name__)


class MessagePriority(Enum):
    """Message priority levels"""
    LOW = 0
    NORMAL = 1
    HIGH = 2
    CRITICAL = 3


class MessageType(Enum):
    """Standard message types"""
    REQUEST = "request"
    RESPONSE = "response"
    EVENT = "event"
    BROADCAST = "broadcast"
    HEARTBEAT = "heartbeat"


@dataclass
class AgentMessage:
    """Represents a message between agents"""
    sender: str
    receiver: str
    message_type: str
    payload: Dict[str, Any]
    priority: MessagePriority = MessagePriority.NORMAL
    correlation_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    reply_to: Optional[str] = None
    expires_at: Optional[datetime] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "sender": self.sender,
            "receiver": self.receiver,
            "message_type": self.message_type,
            "payload": self.payload,
            "priority": self.priority.value,
            "correlation_id": self.correlation_id,
            "timestamp": self.timestamp.isoformat(),
            "reply_to": self.reply_to,
            "expires_at": self.expires_at.isoformat() if self.expires_at else None,
            "metadata": self.metadata
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "AgentMessage":
        return cls(
            sender=data["sender"],
            receiver=data["receiver"],
            message_type=data["message_type"],
            payload=data["payload"],
            priority=MessagePriority(data.get("priority", 1)),
            correlation_id=data.get("correlation_id", str(uuid.uuid4())),
            timestamp=datetime.fromisoformat(data["timestamp"]) if "timestamp" in data else datetime.now(timezone.utc),
            reply_to=data.get("reply_to"),
            expires_at=datetime.fromisoformat(data["expires_at"]) if data.get("expires_at") else None,
            metadata=data.get("metadata", {})
        )


MessageHandler = Callable[[AgentMessage], Awaitable[Optional[Dict[str, Any]]]]


class MessageBus:
    """
    Inter-agent communication via Redis Pub/Sub.

    Features:
    - Publish/Subscribe pattern
    - Request/Response with timeout
    - Priority queues
    - Message persistence for audit
    """

    def __init__(
        self,
        redis_client=None,
        db_pool=None,
        namespace: str = "breakpilot",
        persist_messages: bool = True
    ):
        """
        Initialize the message bus.

        Args:
            redis_client: Async Redis/Valkey client
            db_pool: Async PostgreSQL pool for persistence
            namespace: Channel namespace
            persist_messages: Whether to persist messages for audit
        """
        self.redis = redis_client
        self.db_pool = db_pool
        self.namespace = namespace
        self.persist_messages = persist_messages

        self._handlers: Dict[str, List[MessageHandler]] = {}
        self._pending_responses: Dict[str, asyncio.Future] = {}
        self._subscriptions: Dict[str, asyncio.Task] = {}
        self._running = False

    def _channel(self, agent_id: str) -> str:
        """Generate channel name for agent"""
        return f"{self.namespace}:agent:{agent_id}"

    def _broadcast_channel(self) -> str:
        """Generate broadcast channel name"""
        return f"{self.namespace}:broadcast"

    async def start(self) -> None:
        """Starts the message bus"""
        self._running = True
        logger.info("Message bus started")

    async def stop(self) -> None:
        """Stops the message bus and all subscriptions"""
        self._running = False

        # Cancel all subscription tasks
        for task in self._subscriptions.values():
            task.cancel()

        # Wait for cancellation
        if self._subscriptions:
            await asyncio.gather(
                *self._subscriptions.values(),
                return_exceptions=True
            )

        self._subscriptions.clear()
        logger.info("Message bus stopped")

    async def subscribe(
        self,
        agent_id: str,
        handler: MessageHandler
    ) -> None:
        """
        Subscribe an agent to receive messages.

        Args:
            agent_id: The agent ID to subscribe
            handler: Async function to handle incoming messages
        """
        if agent_id in self._subscriptions:
            logger.warning(f"Agent {agent_id} already subscribed")
            return

        if agent_id not in self._handlers:
            self._handlers[agent_id] = []
        self._handlers[agent_id].append(handler)

        if self.redis:
            # Start Redis subscription
            task = asyncio.create_task(
                self._subscription_loop(agent_id)
            )
            self._subscriptions[agent_id] = task

        logger.info(f"Agent {agent_id} subscribed to message bus")

    async def unsubscribe(self, agent_id: str) -> None:
        """
        Unsubscribe an agent from messages.

        Args:
            agent_id: The agent ID to unsubscribe
        """
        if agent_id in self._subscriptions:
            self._subscriptions[agent_id].cancel()
            try:
                await self._subscriptions[agent_id]
            except asyncio.CancelledError:
                pass
            del self._subscriptions[agent_id]

        self._handlers.pop(agent_id, None)
        logger.info(f"Agent {agent_id} unsubscribed from message bus")

    async def _subscription_loop(self, agent_id: str) -> None:
        """Main subscription loop for an agent"""
        if not self.redis:
            return

        channel = self._channel(agent_id)
        broadcast = self._broadcast_channel()

        pubsub = self.redis.pubsub()
        await pubsub.subscribe(channel, broadcast)

        try:
            while self._running:
                message = await pubsub.get_message(
                    ignore_subscribe_messages=True,
                    timeout=1.0
                )

                if message and message["type"] == "message":
                    await self._handle_incoming_message(
                        agent_id,
                        message["data"]
                    )
        except asyncio.CancelledError:
            pass
        finally:
            await pubsub.unsubscribe(channel, broadcast)
            await pubsub.close()

    async def _handle_incoming_message(
        self,
        agent_id: str,
        raw_data: bytes
    ) -> None:
        """Process an incoming message"""
        try:
            data = json.loads(raw_data)
            message = AgentMessage.from_dict(data)

            # Check if this is a response to a pending request
            if message.correlation_id in self._pending_responses:
                future = self._pending_responses[message.correlation_id]
                if not future.done():
                    future.set_result(message.payload)
                return

            # Call handlers
            handlers = self._handlers.get(agent_id, [])
            for handler in handlers:
                try:
                    response = await handler(message)

                    # If handler returns data and there's a reply_to, send response
                    if response and message.reply_to:
                        await self.publish(AgentMessage(
                            sender=agent_id,
                            receiver=message.sender,
                            message_type="response",
                            payload=response,
                            correlation_id=message.correlation_id
                        ))

                except Exception as e:
                    logger.error(
                        f"Error in message handler for {agent_id}: {e}"
                    )

        except Exception as e:
            logger.error(f"Error processing message: {e}")

    async def publish(self, message: AgentMessage) -> None:
        """
        Publishes a message to an agent.

        Args:
            message: The message to publish
        """
        # Persist message if enabled
        if self.persist_messages:
            await self._persist_message(message)

        if self.redis:
            channel = self._channel(message.receiver)
            await self.redis.publish(
                channel,
                json.dumps(message.to_dict())
            )
        else:
            # Local delivery for testing
            await self._local_deliver(message)

        logger.debug(
            f"Published message from {message.sender} to {message.receiver}: "
            f"{message.message_type}"
        )

    async def broadcast(self, message: AgentMessage) -> None:
        """
        Broadcasts a message to all agents.

        Args:
            message: The message to broadcast
        """
        message.receiver = "*"  # Indicate broadcast

        if self.persist_messages:
            await self._persist_message(message)

        if self.redis:
            await self.redis.publish(
                self._broadcast_channel(),
                json.dumps(message.to_dict())
            )
        else:
            # Local broadcast
            for agent_id in self._handlers:
                await self._local_deliver(message, agent_id)

        logger.debug(f"Broadcast message from {message.sender}: {message.message_type}")

    async def request(
        self,
        message: AgentMessage,
        timeout: float = 30.0
    ) -> Dict[str, Any]:
        """
        Sends a request and waits for a response.

        Args:
            message: The request message
            timeout: Timeout in seconds

        Returns:
            Response payload

        Raises:
            TimeoutError: If no response within timeout
        """
        # Mark this as a request that needs a response
        message.reply_to = message.sender

        # Create future for response
        future: asyncio.Future = asyncio.Future()
        self._pending_responses[message.correlation_id] = future

        try:
            # Publish as task so handler runs concurrently (enables timeout)
            publish_task = asyncio.create_task(self.publish(message))

            # Wait for response with timeout
            try:
                return await asyncio.wait_for(future, timeout)
            finally:
                if not publish_task.done():
                    publish_task.cancel()
                    try:
                        await publish_task
                    except (asyncio.CancelledError, Exception):
                        pass

        except asyncio.TimeoutError:
            logger.warning(
                f"Request timeout: {message.sender} -> {message.receiver} "
                f"({message.message_type})"
            )
            raise

        finally:
            # Clean up
            self._pending_responses.pop(message.correlation_id, None)

    async def _local_deliver(
        self,
        message: AgentMessage,
        target_agent: Optional[str] = None
    ) -> None:
        """Local message delivery for testing without Redis"""
        agent_id = target_agent or message.receiver
        handlers = self._handlers.get(agent_id, [])

        for handler in handlers:
            try:
                response = await handler(message)
                if response and message.reply_to:
                    if message.correlation_id in self._pending_responses:
                        future = self._pending_responses[message.correlation_id]
                        if not future.done():
                            future.set_result(response)
            except Exception as e:
                logger.error(f"Error in local handler: {e}")

    async def _persist_message(self, message: AgentMessage) -> None:
        """Persist message to PostgreSQL for audit"""
        if not self.db_pool:
            return

        try:
            async with self.db_pool.acquire() as conn:
                await conn.execute(
                    """
                    INSERT INTO agent_messages
                        (id, sender, receiver, message_type, payload,
                         priority, correlation_id, created_at)
                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                    """,
                    str(uuid.uuid4()),
                    message.sender,
                    message.receiver,
                    message.message_type,
                    json.dumps(message.payload),
                    message.priority.value,
                    message.correlation_id,
                    message.timestamp
                )
        except Exception as e:
            logger.warning(f"Failed to persist message: {e}")

    def on_message(self, message_type: str):
        """
        Decorator for message handlers.

        Usage:
            @bus.on_message("grade_request")
            async def handle_grade(message):
                return {"score": 85}
        """
        def decorator(func: MessageHandler):
            async def wrapper(message: AgentMessage):
                if message.message_type == message_type:
                    return await func(message)
                return None
            return wrapper
        return decorator

    async def get_message_history(
        self,
        agent_id: Optional[str] = None,
        message_type: Optional[str] = None,
        limit: int = 100
    ) -> List[Dict[str, Any]]:
        """
        Gets message history from persistence.

        Args:
            agent_id: Filter by sender or receiver
            message_type: Filter by message type
            limit: Maximum results

        Returns:
            List of message dicts
        """
        if not self.db_pool:
            return []

        query = """
            SELECT sender, receiver, message_type, payload, priority,
                   correlation_id, created_at
            FROM agent_messages
            WHERE 1=1
        """
        params = []

        if agent_id:
            query += " AND (sender = $1 OR receiver = $1)"
            params.append(agent_id)

        if message_type:
            param_num = len(params) + 1
            query += f" AND message_type = ${param_num}"
            params.append(message_type)

        query += f" ORDER BY created_at DESC LIMIT {limit}"

        async with self.db_pool.acquire() as conn:
            rows = await conn.fetch(query, *params)
            return [dict(row) for row in rows]

    @property
    def connected(self) -> bool:
        """Returns whether the bus is connected to Redis"""
        return self.redis is not None and self._running

    @property
    def subscriber_count(self) -> int:
        """Returns number of subscribed agents"""
        return len(self._subscriptions)
