main.py

from contextlib import asynccontextmanager
from typing import List, Optional
from uuid import uuid4

import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Depends, status
from fastapi.middleware.cors import CORSMiddleware

from core.manager.websocket_manager import WebSocketManager
from core.settings.settings import Settings
from core.types.datamodels import (
    ClientInfo,
    SubscriptionRequest,
    Subscription,
    NotificationRequest,
    Notification,
    BroadcastRequest,
)


class WebSocketNotificationService:
    def __init__(self, settings: Optional[Settings] = None):
        self.settings = settings or Settings()
        self.manager = WebSocketManager(settings=self.settings)
        self.logger = self.manager.logger
        self.app = self._create_app()
        self._configure_middleware()
        self._register_routes()

    def _create_app(self) -> FastAPI:
        return FastAPI(
            title=self.settings.APP_NAME,
            description="Serivce de notification WebSocket et de real-time data",
            version=self.settings.API_VERSION,
            lifespan=self._lifespan,
            docs_url="/docs" if self.settings.APP_DEBUG else None,
            redoc_url="/redoc" if self.settings.APP_DEBUG else None,
            websocket_ping_interval=self.settings.WEBSOCKET_PING_INTERVAL,
            websocket_timeout=self.settings.WEBSOCKET_TIMEOUT,
            openapi_tags=[
                {
                    "name": "health",
                    "description": "Santé et vérification du système"
                },
                {
                    "name": "websocket",
                    "description": "Gestion de connexions WebSocket"
                },
                {
                    "name": "subscriptions",
                    "description": "Gestion des subscriptions"
                },
                {
                    "name": "notifications",
                    "description": "Envoie et gestion des notifications"
                },
                {
                    "name": "monitoring",
                    "description": "Analytics et statistiques"
                }
            ]
        )

    def _configure_middleware(self):
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=[self.settings.WS_CORS_ORIGINS],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

    @asynccontextmanager
    async def _lifespan(self, _app: FastAPI):
        self.logger.info("🚀 Starting WebSocket notification service")

        try:
            await self.manager.initialize()
            yield
        finally:
            await self.manager.shutdown()
            self.logger.info("🛑 WebSocket notification service stopped")

    def _register_routes(self):
        self._register_health_routes()
        self._register_websocket_routes()
        self._register_subscription_routes()
        self._register_notification_routes()
        self._register_monitoring_routes()

    def _register_health_routes(self):
        @self.app.get("/", tags=["health"])
        async def root():
            return {
                "service": self.settings.APP_NAME,
                "version": self.settings.API_VERSION,
                "status": "healthy",
                "timestamp": self.manager.get_current_timestamp()
            }

        @self.app.get("/health", tags=["health"])
        async def health_check():
            return {
                "status": "healthy",
                "connections": len(self.manager.active_connections),
                "uptime": self.manager.get_uptime(),
                "memory_usage": self.manager.get_memory_usage()
            }

    def _register_websocket_routes(self):
        @self.app.websocket("/ws/{user_id}")
        async def websocket_endpoint(websocket: WebSocket, user_id: str):
            session_id = ""

            try:
                client_info = ClientInfo(
                    session_id=str(uuid4()),
                    user_id=user_id,
                    user_roles=self._extract_roles(websocket),
                    ip_address=self._get_client_ip(websocket),
                    user_agent=self._get_user_agent(websocket)
                )

                session_id = await self.manager.connect(websocket, client_info)

                await self._handle_websocket_messages(websocket, session_id)

            except WebSocketDisconnect:
                self.logger.info(f"WebSocket connection closed: {session_id}")
            except Exception as e:
                self.logger.error(f"WebSocket error: {e}")
                await self._handle_websocket_error(websocket, e)
            finally:
                self.manager.disconnect(session_id)

    def _register_subscription_routes(self):
        @self.app.post("/subscriptions/{session_id}", tags=["subscriptions"])
        async def add_subscription(
                subscription_request: SubscriptionRequest,
                session_id: str = Depends(self._validate_session),
        ):
            try:
                subscription = Subscription(**subscription_request.model_dump())

                if not self.manager.add_subscription(session_id, subscription):
                    raise HTTPException(
                        status_code=status.HTTP_404_NOT_FOUND,
                        detail="Failed to add subscription"
                    )

                return {
                    "success": True,
                    "message": "Subscription added successfully",
                    "subscription_id": subscription.id,
                    "timestamp": self.manager.get_current_timestamp()
                }
            except Exception as e:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=str(e)
                )

        @self.app.delete("/subscriptions/{session_id}/{subscription_id}", tags=["subscriptions"])
        async def remove_subscription(
                subscription_id: str,
                session_id: str = Depends(self._validate_session),
        ):
            if not self.manager.remove_subscription(session_id, subscription_id):
                raise HTTPException(
                    status_code=status.HTTP_404_NOT_FOUND,
                    detail="Subscription not found"
                )

            return {
                "success": True,
                "message": "Subscription removed successfully",
                "timestamp": self.manager.get_current_timestamp()
            }

        @self.app.get("/subscriptions/{session_id}", tags=["subscriptions"])
        async def get_subscriptions(
                session_id: str = Depends(self._validate_session),
        ):
            """Get all subscriptions for a session"""
            subscriptions = self.manager.get_subscriptions(session_id)
            return {
                "success": True,
                "subscriptions": subscriptions,
                "count": len(subscriptions),
                "timestamp": self.manager.get_current_timestamp()
            }

    def _register_notification_routes(self):
        @self.app.post("/notifications/user/{user_id}", tags=["notifications"])
        async def send_user_notification(
                user_id: str,
                notification_request: NotificationRequest,
        ):
            try:
                notification = Notification(**notification_request.model_dump())
                result = await self.manager.send_to_user(user_id, notification)

                return {
                    "success": True,
                    "message": f"Notification sent to user {user_id}",
                    "notification_id": notification.id,
                    "delivered_to": result.get("delivered_to", 0),
                    "timestamp": self.manager.get_current_timestamp()
                }
            except Exception as e:
                raise HTTPException(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    detail=str(e)
                )

        @self.app.post("/notifications/role/{role}", tags=["notifications"])
        async def send_role_notification(
                role: str,
                notification_request: NotificationRequest,
        ):
            try:
                notification = Notification(**notification_request.model_dump())
                result = await self.manager.send_to_role(role, notification)

                return {
                    "success": True,
                    "message": f"Notification sent to role {role}",
                    "notification_id": notification.id,
                    "delivered_to": result.get("delivered_to", 0),
                    "timestamp": self.manager.get_current_timestamp()
                }
            except Exception as e:
                raise HTTPException(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    detail=str(e)
                )

        @self.app.post("/notifications/broadcast", tags=["notifications"])
        async def broadcast_notification(
                notification_request: NotificationRequest,
        ):
            try:
                notification = Notification(**notification_request.model_dump())
                result = await self.manager.broadcast_notification(notification)

                return {
                    "success": True,
                    "message": "Notification broadcasted",
                    "notification_id": notification.id,
                    "delivered_to": result.get("delivered_to", 0),
                    "timestamp": self.manager.get_current_timestamp()
                }
            except Exception as e:
                raise HTTPException(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    detail=str(e)
                )

        @self.app.post("/notifications/send", tags=["notifications"])
        async def send_targeted_notification(
                request: BroadcastRequest,
        ):
            try:
                notification = Notification(**request.notification.model_dump())

                if request.target_type == "user":
                    result = await self.manager.send_to_user(request.target_value, notification)
                elif request.target_type == "role":
                    result = await self.manager.send_to_role(request.target_value, notification)
                elif request.target_type == "global":
                    result = await self.manager.broadcast_notification(notification)
                else:
                    raise ValueError("Invalid target type")

                return {
                    "success": True,
                    "message": f"Notification sent to {request.target_type}: {request.target_value}",
                    "notification_id": notification.id,
                    "delivered_to": result.get("delivered_to", 0),
                    "timestamp": self.manager.get_current_timestamp()
                }
            except Exception as e:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=str(e)
                )

    def _register_monitoring_routes(self):
        @self.app.get("/notifications/history", tags=["monitoring"])
        async def get_notification_history(
                limit: int = 100,
                offset: int = 0,
        ):
            history = self.manager.get_notification_history(limit, offset)
            return {
                "success": True,
                "history": history,
                "count": len(history),
                "limit": limit,
                "offset": offset,
                "timestamp": self.manager.get_current_timestamp()
            }

        @self.app.get("/stats", tags=["monitoring"])
        async def get_service_stats():
            stats = self.manager.get_stats()
            return {
                "success": True,
                "stats": stats,
                "timestamp": self.manager.get_current_timestamp()
            }

        @self.app.get("/clients", tags=["monitoring"])
        async def get_connected_clients():
            clients = [
                {
                    "session_id": session_id,
                    "user_id": info.user_id,
                    "user_roles": info.user_roles,
                    "ip_address": info.ip_address,
                    "connected_at": info.connected_at.isoformat(),
                    "subscriptions_count": len(self.manager.subscriptions.get(session_id, [])),
                    "last_activity": info.last_activity.isoformat() if hasattr(info, 'last_activity') else None
                }
                for session_id, info in self.manager.client_info.items()
            ]

            return {
                "success": True,
                "clients": clients,
                "total_connections": len(clients),
                "timestamp": self.manager.get_current_timestamp()
            }

    # Helper methods
    async def _validate_session(self, session_id: str) -> str:
        if session_id not in self.manager.active_connections:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="Session not found or already closed"
            )
        return session_id

    @staticmethod
    def _extract_roles(websocket: WebSocket) -> List[str]:
        roles_param = websocket.query_params.get("roles", "")
        return roles_param.split(",") if roles_param else []

    @staticmethod
    def _get_client_ip(websocket: WebSocket) -> str:
        return websocket.client.host if websocket.client else "unknown"

    @staticmethod
    def _get_user_agent(websocket: WebSocket) -> str:
        return websocket.headers.get("user-agent", "unknown")

    async def _handle_websocket_messages(self, websocket: WebSocket, session_id: str):
        while True:
            try:
                message = await websocket.receive_text()
                await self._process_websocket_message(websocket, session_id, message)
            except WebSocketDisconnect:
                raise
            except Exception as e:
                self.logger.error(f"Error processing WebSocket message: {e}")
                await websocket.send_text(f"Error: {str(e)}")

    async def _process_websocket_message(self, websocket: WebSocket, session_id: str, message: str):
        if message == "ping":
            await websocket.send_text("pong")
        elif message.startswith("subscribe:"):
            topic = message.replace("subscribe:", "")
            subscription = Subscription(type="topic", target=topic)
            self.manager.add_subscription(session_id, subscription)
            await websocket.send_text(f"subscribed:{topic}")
        elif message.startswith("unsubscribe:"):
            topic = message.replace("unsubscribe:", "")
            await websocket.send_text(f"unsubscribed:{topic}")
        else:
            await self._handle_custom_message(websocket, session_id, message)

    async def _handle_custom_message(self, websocket: WebSocket, session_id: str, message: str):
        """Handle custom WebSocket messages"""
        pass

    @staticmethod
    async def _handle_websocket_error(websocket: WebSocket, error: Exception):
        """Handle WebSocket errors gracefully"""
        try:
            await websocket.send_text(f"Error: {str(error)}")
        except:
            raise

    def run(self):
        uvicorn.run(
            self.app,
            host=self.settings.APP_HOST,
            port=self.settings.APP_PORT,
            reload=self.settings.APP_DEBUG,
            log_level="info",
            server_header=False,
        )


def create_app(settings: Optional[Settings] = None) -> FastAPI:
    service = WebSocketNotificationService(settings)
    return service.app


def main():
    service = WebSocketNotificationService()
    service.run()


app = create_app()

if __name__ == "__main__":
    main()