core/manager/websocket_manager.py

import asyncio
from collections import deque
from datetime import datetime, timezone
from typing import Dict, List, Any, Deque, Optional

import psutil
from fastapi import WebSocket

from core.settings.settings import Settings
from core.types.datamodels import ClientInfo, Subscription, Notification
from core.utils.utils import current_utc_iso


class WebSocketManager:
    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self.logger = self.settings.logger

        self.active_connections: Dict[str, WebSocket] = {}
        self.client_info: Dict[str, ClientInfo] = {}
        self.subscriptions: Dict[str, List[Subscription]] = {}
        self.notification_history: Deque[Notification] = deque(
            maxlen=self.settings.get("MAX_NOTIFICATION_HISTORY", 1000)
        )

        self._start_time: Optional[datetime] = None
        self._is_initialized: bool = False
        self._shutdown_event: Optional[asyncio.Event] = None
        self._background_tasks: List[asyncio.Task] = []

        self._total_messages_sent: int = 0
        self._total_connections_handled: int = 0
        self._connection_errors: int = 0

    async def initialize(self) -> None:
        if self._is_initialized:
            self.logger.warning("WebSocketManager already initialized")
            return

        self._start_time = datetime.now(timezone.utc)
        self._shutdown_event = asyncio.Event()
        self._is_initialized = True

        self._background_tasks = [
            asyncio.create_task(self._cleanup_expired_notifications()),
            asyncio.create_task(self._heartbeat_checker()),
            asyncio.create_task(self._stats_collector())
        ]

        self.logger.info("WebSocketManager initialized successfully")

    async def shutdown(self) -> None:
        if not self._is_initialized:
            return

        self.logger.info("Shutting down WebSocketManager...")

        if self._shutdown_event:
            self._shutdown_event.set()

        for task in self._background_tasks:
            if not task.done():
                task.cancel()

        if self._background_tasks:
            await asyncio.gather(*self._background_tasks, return_exceptions=True)

        await self._close_all_connections()

        self.active_connections.clear()
        self.client_info.clear()
        self.subscriptions.clear()
        self.notification_history.clear()

        self._is_initialized = False
        self.logger.info("WebSocketManager shutdown completed")

    async def connect(self, websocket: WebSocket, client: ClientInfo) -> str:
        await websocket.accept()
        session_id = client.session_id

        self.active_connections[session_id] = websocket
        self.client_info[session_id] = client
        self.subscriptions[session_id] = []

        self._total_connections_handled += 1

        self.logger.info(f"Nouvelle connexion: {session_id}")

        await self._send(
            websocket,
            {
                "session_id": session_id,
                "type": "connection_established",
                "title": "Connexion réussie",
                "message": "Connexion au service de notifications établie",
                "timestamp": current_utc_iso(),
            },
        )

        return session_id

    def disconnect(self, session_id: str) -> None:
        if session_id in self.active_connections:
            del self.active_connections[session_id]
            del self.client_info[session_id]
            del self.subscriptions[session_id]
            self.logger.info(f"Connexion fermée: {session_id}")

    async def send_personal_message(
            self, session_id: str, payload: Dict[str, Any]
    ) -> None:
        if ws := self.active_connections.get(session_id):
            try:
                await self._send(ws, payload)
                self._total_messages_sent += 1
            except RuntimeError as e:
                self.logger.error(f"Erreur d'envoi à {session_id}: {e}")
                self._connection_errors += 1
                self.disconnect(session_id)

    async def broadcast_notification(self, notification: Notification) -> Dict[str, Any]:
        if notification.is_expired:
            return {"delivered_to": 0}

        self._add_to_history(notification)
        payload = self._format_notification(notification)
        delivered_count = 0

        for session_id, subscriptions in self.subscriptions.items():
            if not (client := self.client_info.get(session_id)):
                continue

            if self._should_receive(subscriptions, notification, client):
                await self.send_personal_message(session_id, payload)
                delivered_count += 1

        return {"delivered_to": delivered_count}

    async def send_to_user(self, user_id: str, notification: Notification) -> Dict[str, Any]:
        if notification.is_expired:
            return {"delivered_to": 0}

        self._add_to_history(notification)
        payload = self._format_notification(notification)
        delivered_count = 0

        for session_id, client in self.client_info.items():
            if client.user_id == user_id:
                await self.send_personal_message(session_id, payload)
                delivered_count += 1

        return {"delivered_to": delivered_count}

    async def send_to_role(self, role: str, notification: Notification) -> Dict[str, Any]:
        if notification.is_expired:
            return {"delivered_to": 0}

        self._add_to_history(notification)
        payload = self._format_notification(notification)
        delivered_count = 0

        for session_id, client in self.client_info.items():
            if role in client.user_roles:
                await self.send_personal_message(session_id, payload)
                delivered_count += 1

        return {"delivered_to": delivered_count}

    def add_subscription(self, session_id: str, subscription: Subscription) -> bool:
        if not (sub_list := self.subscriptions.get(session_id)):
            return False

        for existing in sub_list:
            if existing.id == subscription.id:
                return False

        sub_list.append(subscription)
        return True

    def update_subscription(self, session_id: str, subscription: Subscription) -> bool:
        if not (sub_list := self.subscriptions.get(session_id)):
            return False

        for i, existing in enumerate(sub_list):
            if existing.id == subscription.id:
                sub_list[i] = subscription
                return True
        return False

    def remove_subscription(self, session_id: str, subscription_id: str) -> bool:
        if not (sub_list := self.subscriptions.get(session_id)):
            return False

        initial_count = len(sub_list)
        self.subscriptions[session_id] = [
            sub for sub in sub_list if sub.id != subscription_id
        ]
        return len(self.subscriptions[session_id]) < initial_count

    def get_subscriptions(self, session_id: str) -> List[Subscription]:
        return self.subscriptions.get(session_id, [])[:]

    def get_notification_history(self, limit: int = 100, offset: int = 0) -> List[Notification]:
        history = [n for n in list(self.notification_history) if not n.is_expired]
        return history[offset:offset + limit]

    def get_stats(self) -> Dict[str, Any]:
        memory_info = self.get_memory_usage()

        return {
            "active_connections": len(self.active_connections),
            "total_subscriptions": sum(len(s) for s in self.subscriptions.values()),
            "total_notifications": len(self.notification_history),
            "total_messages_sent": self._total_messages_sent,
            "total_connections_handled": self._total_connections_handled,
            "connection_errors": self._connection_errors,
            "uptime_seconds": self.get_uptime(),
            "memory_usage": memory_info,
            "timestamp": current_utc_iso(),
        }

    @staticmethod
    def get_current_timestamp() -> str:
        return current_utc_iso()

    def get_uptime(self) -> str:
        if not self._start_time:
            return self.format_uptime(0)
        return self.format_uptime((datetime.now(timezone.utc) - self._start_time).total_seconds())

    def get_memory_usage(self) -> Dict[str, Any]:
        try:
            process = psutil.Process()
            memory_info = process.memory_info()

            return {
                "resident_set_size_mb": self._format_size(memory_info.rss),
                "virtual_memory_size_mb": self._format_size(memory_info.vms),
                "percent": self._format_size(process.memory_percent()),
                "available_mb": self._format_size(psutil.virtual_memory().available),
                "total_mb": self._format_size(psutil.virtual_memory().total),
            }
        except Exception as e:
            self.logger.error(f"Error getting memory usage: {e}")
            return {
                "resident_set_size_mb": 0,
                "virtual_memory_size_mb": 0,
                "percent": 0,
                "available_mb": 0,
                "total_mb": 0,
                "error": str(e)
            }

    async def _cleanup_expired_notifications(self) -> None:
        while not self._shutdown_event.is_set():
            try:
                active_notifications = [
                    n for n in self.notification_history if not n.is_expired
                ]

                self.notification_history.clear()
                self.notification_history.extend(active_notifications)

                await asyncio.sleep(self.settings.get("CLEANUP_INTERVAL", 300))  # 5 minutes

            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"Error in cleanup task: {e}")
                await asyncio.sleep(60)

    async def _heartbeat_checker(self) -> None:
        while not self._shutdown_event.is_set():
            try:
                disconnected_sessions = []

                for session_id, websocket in self.active_connections.items():
                    try:
                        await websocket.send_json(
                            {"session_id": session_id, "type": "ping", "timestamp": current_utc_iso()})
                    except Exception as e:
                        self.logger.warning(f"Connection {session_id} appears to be dead: {e}")
                        disconnected_sessions.append(session_id)

                for session_id in disconnected_sessions:
                    self.disconnect(session_id)

                await asyncio.sleep(self.settings.get("HEARTBEAT_INTERVAL", 30))  # 30 seconds

            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"Error in heartbeat checker: {e}")
                await asyncio.sleep(60)

    async def _stats_collector(self) -> None:
        while not self._shutdown_event.is_set():
            try:
                _ = self.get_stats()

                await asyncio.sleep(self.settings.get("STATS_INTERVAL", 60))  # 1 minute

            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"Error in stats collector: {e}")
                await asyncio.sleep(60)

    async def _close_all_connections(self) -> None:
        if not self.active_connections:
            return

        self.logger.info(f"Closing {len(self.active_connections)} active connections...")

        shutdown_message = {
            "type": "service_shutdown",
            "title": "Service en cours d'arrêt",
            "message": "Le service de notifications va être arrêté",
            "timestamp": current_utc_iso()
        }

        close_tasks = []
        for session_id, websocket in self.active_connections.items():
            try:
                await self._send(websocket, shutdown_message)
                close_tasks.append(websocket.close())
            except Exception as e:
                self.logger.error(f"Error closing connection {session_id}: {e}")

        if close_tasks:
            await asyncio.gather(*close_tasks, return_exceptions=True)

    async def _send(self, websocket: WebSocket, payload: Dict[str, Any]) -> None:
        try:
            await websocket.send_json(payload)
        except Exception as e:
            self.logger.error(f"Failed to send message: {e}")
            raise RuntimeError(f"WebSocket send failed: {e}")

    def _add_to_history(self, notification: Notification) -> None:
        if not notification.is_expired:
            self.notification_history.append(notification)

    @staticmethod
    def _should_receive(
            subscriptions: List[Subscription],
            notification: Notification,
            client: ClientInfo,
    ) -> bool:
        return any(
            sub.matches(notification, client.user_id, client.user_roles)
            for sub in subscriptions
        )

    @staticmethod
    def _format_notification(notification: Notification) -> Dict[str, Any]:
        return {"type": "notification", **notification.model_dump(exclude_unset=True)}

    @staticmethod
    def _format_size(bytes_value) -> str:
        if bytes_value < 1024:
            return f"{bytes_value} B"
        elif bytes_value < 1024 ** 2:
            return f"{bytes_value / 1024:.2f} KB"
        elif bytes_value < 1024 ** 3:
            return f"{bytes_value / (1024 ** 2):.2f} MB"
        else:
            return f"{bytes_value / (1024 ** 3):.2f} GB"

    @staticmethod
    def format_uptime(seconds: float) -> str:
        minutes, sec = divmod(int(seconds), 60)
        hours, minutes = divmod(minutes, 60)
        days, hours = divmod(hours, 24)
        parts = []
        if days: parts.append(f"{days}d")
        if hours: parts.append(f"{hours}h")
        if minutes: parts.append(f"{minutes}m")
        if sec or not parts: parts.append(f"{sec}s")
        return " ".join(parts)