import asyncio
from typing import Dict, List, Any, Optional
import json
from datetime import datetime, timezone

class NotificationManager:
    def __init__(self):
        # Dictionary of active queues: {user_id: [Queue, Queue, ...]}
        # We use a list of queues per user to support multiple tabs/sessions
        self.active_connections: Dict[int, List[asyncio.Queue]] = {}
        self.lock = asyncio.Lock()

    async def subscribe(self, user_id: int) -> asyncio.Queue:
        """Create a new queue for a user's SSE connection."""
        queue = asyncio.Queue()
        async with self.lock:
            if user_id not in self.active_connections:
                self.active_connections[user_id] = []
            self.active_connections[user_id].append(queue)
        return queue

    async def unsubscribe(self, user_id: int, queue: asyncio.Queue):
        """Remove a queue when an SSE connection closes."""
        async with self.lock:
            if user_id in self.active_connections:
                if queue in self.active_connections[user_id]:
                    self.active_connections[user_id].remove(queue)
                if not self.active_connections[user_id]:
                    self.active_connections.pop(user_id, None)

    async def broadcast(self, notification_data: Dict[str, Any], target_user_id: Optional[int] = None):
        """
        Send a notification to active subscribers.
        If target_user_id is None, broadcast to ALL admins (currently all admins see all notifications).
        """
        # Prepare the message
        # Convert datetime to string if present
        if isinstance(notification_data.get("created_at"), datetime):
            notification_data["created_at"] = notification_data["created_at"].isoformat()
        
        message = json.dumps(notification_data)
        
        async with self.lock:
            if target_user_id:
                # Send to specific user
                if target_user_id in self.active_connections:
                    for queue in self.active_connections[target_user_id]:
                        await queue.put(message)
            else:
                # Broadcast to everyone (Admins)
                for user_queues in self.active_connections.values():
                    for queue in user_queues:
                        await queue.put(message)

# Global Manager Instance
notification_manager = NotificationManager()