ai-voicebot/server/main_refactored.py

"""
Refactored main.py - Step 1 of Server Architecture Improvement

This is a refactored version of the original main.py that demonstrates the new
modular architecture with separated concerns:

- SessionManager: Handles session lifecycle and persistence
- LobbyManager: Handles lobby management and chat
- AuthManager: Handles authentication and name protection
- WebSocket message routing: Clean message handling
- Separated API modules: Admin, session, and lobby endpoints

This maintains backward compatibility while providing a foundation for
further improvements.
"""

from __future__ import annotations
import os
from contextlib import asynccontextmanager

from fastapi import FastAPI, WebSocket, Path
from fastapi.staticfiles import StaticFiles

# Import our new modular components
try:
    from core.session_manager import SessionManager
    from core.lobby_manager import LobbyManager
    from core.auth_manager import AuthManager
    from websocket.connection import WebSocketConnectionManager
    from api.admin import AdminAPI
    from api.sessions import SessionAPI
    from api.lobbies import LobbyAPI
except ImportError:
    # Handle relative imports when running as module
    import sys
    import os

    sys.path.append(os.path.dirname(os.path.abspath(__file__)))

    from core.session_manager import SessionManager
    from core.lobby_manager import LobbyManager
    from core.auth_manager import AuthManager
    from websocket.connection import WebSocketConnectionManager
    from api.admin import AdminAPI
    from api.sessions import SessionAPI
    from api.lobbies import LobbyAPI

from logger import logger


# Configuration
public_url = os.getenv("PUBLIC_URL", "/")
if not public_url.endswith("/"):
    public_url += "/"

ADMIN_TOKEN = os.getenv("ADMIN_TOKEN", None)

# Global managers - these replace the global variables from original main.py
session_manager: SessionManager = None
lobby_manager: LobbyManager = None
auth_manager: AuthManager = None
websocket_manager: WebSocketConnectionManager = None

# API routers
admin_api: AdminAPI = None
session_api: SessionAPI = None
lobby_api: LobbyAPI = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Lifespan context manager for startup and shutdown events"""
    global session_manager, lobby_manager, auth_manager, websocket_manager
    global admin_api, session_api, lobby_api

    # Startup
    logger.info("Starting AI Voice Bot server with modular architecture...")

    # Initialize managers
    session_manager = SessionManager("sessions.json")
    lobby_manager = LobbyManager()
    auth_manager = AuthManager("sessions.json")

    # Load existing data
    session_manager.load()

    # Restore lobbies for existing sessions
    # Note: This is a simplified version - full lobby restoration would be more complex
    for session in session_manager.get_all_sessions():
        for lobby_info in session.lobbies:
            # Create lobby if it doesn't exist
            lobby = lobby_manager.create_or_get_lobby(
                name=lobby_info.name, private=lobby_info.private
            )
            # Add session to lobby (but don't trigger events during startup)
            with lobby.lock:
                lobby.sessions[session.id] = session

    # Set up dependency injection for name protection
    lobby_manager.set_name_protection_checker(auth_manager.is_name_protected)

    # Initialize WebSocket manager
    websocket_manager = WebSocketConnectionManager(
        session_manager=session_manager,
        lobby_manager=lobby_manager,
        auth_manager=auth_manager,
    )

    # Initialize API routers
    admin_api = AdminAPI(
        session_manager=session_manager,
        lobby_manager=lobby_manager,
        auth_manager=auth_manager,
        admin_token=ADMIN_TOKEN,
        public_url=public_url,
    )

    session_api = SessionAPI(session_manager=session_manager, public_url=public_url)

    lobby_api = LobbyAPI(
        session_manager=session_manager,
        lobby_manager=lobby_manager,
        public_url=public_url,
    )

    # Register API routes
    app.include_router(admin_api.router)
    app.include_router(session_api.router)
    app.include_router(lobby_api.router)

    # Start background tasks
    await session_manager.start_background_tasks()

    logger.info("AI Voice Bot server started successfully!")
    logger.info(f"Server URL: {public_url}")
    logger.info(f"Sessions loaded: {session_manager.get_session_count()}")
    logger.info(f"Lobbies available: {lobby_manager.get_lobby_count()}")
    logger.info(f"Protected names: {auth_manager.get_protection_count()}")

    if ADMIN_TOKEN:
        logger.info("Admin endpoints protected with token")
    else:
        logger.warning("Admin endpoints are unprotected")

    yield

    # Shutdown
    logger.info("Shutting down AI Voice Bot server...")

    # Stop background tasks
    if session_manager:
        await session_manager.stop_background_tasks()

    logger.info("Server shutdown complete")


# Create FastAPI app
app = FastAPI(
    title="AI Voice Bot Server (Refactored)",
    description="WebRTC voice chat server with modular architecture",
    version="2.0.0",
    lifespan=lifespan,
)

logger.info(f"Starting server with public URL: {public_url}")


@app.websocket(f"{public_url}" + "ws/lobby/{lobby_id}/{session_id}")
async def lobby_websocket(
    websocket: WebSocket,
    lobby_id: str | None = Path(...),
    session_id: str | None = Path(...),
):
    """WebSocket endpoint for lobby connections - now uses WebSocketConnectionManager"""
    await websocket_manager.handle_connection(websocket, lobby_id, session_id)


# Serve static files if available (for client)
try:
    app.mount(public_url + "static", StaticFiles(directory="static"), name="static")
    logger.info("Static files mounted at /static")
except Exception:
    logger.info("No static directory found, skipping static file serving")


# Health check for the new architecture
@app.get(f"{public_url}api/system/health")
def system_health():
    """System health check showing manager status"""
    return {
        "status": "ok",
        "architecture": "modular",
        "version": "2.0.0",
        "managers": {
            "session_manager": "active" if session_manager else "inactive",
            "lobby_manager": "active" if lobby_manager else "inactive",
            "auth_manager": "active" if auth_manager else "inactive",
            "websocket_manager": "active" if websocket_manager else "inactive",
        },
        "statistics": {
            "sessions": session_manager.get_session_count() if session_manager else 0,
            "lobbies": lobby_manager.get_lobby_count() if lobby_manager else 0,
            "protected_names": auth_manager.get_protection_count()
            if auth_manager
            else 0,
        },
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)