api.py

import asyncio
import json
import logging
import os
from pathlib import Path
import time
import uuid
from aiohttp import web, WSMsgType
from typing import Dict, Any

from server.api_core import VideoGenerationAPI
from server.api_session import SessionManager
from server.api_metrics import MetricsTracker
from server.api_config import *

# Set up colored logging
from server.logging_utils import setup_colored_logging, get_logger

setup_colored_logging()
logger = get_logger(__name__)

# Create global session and metrics managers
session_manager = SessionManager()
metrics_tracker = MetricsTracker()

# Dictionary to track connected anonymous clients by IP address
anon_connections = {}
anon_connection_lock = asyncio.Lock()

async def status_handler(request: web.Request) -> web.Response:
    """Handler for API status endpoint"""
    api = session_manager.shared_api
    
    # Get current busy status of all endpoints
    endpoint_statuses = []
    for ep in api.endpoint_manager.endpoints:
        endpoint_statuses.append({
            'id': ep.id,
            'url': ep.url,
            'busy': ep.busy,
            'last_used': ep.last_used,
            'error_count': ep.error_count,
            'error_until': ep.error_until
        })
    
    # Get session statistics
    session_stats = session_manager.get_session_stats()
    
    # Get metrics
    api_metrics = metrics_tracker.get_metrics()
    
    return web.json_response({
        'product': PRODUCT_NAME,
        'version': PRODUCT_VERSION,
        'maintenance_mode': MAINTENANCE_MODE,
        'available_endpoints': len(VIDEO_ROUND_ROBIN_ENDPOINT_URLS),
        'endpoint_status': endpoint_statuses,
        'active_endpoints': sum(1 for ep in endpoint_statuses if not ep['busy'] and ('error_until' not in ep or ep['error_until'] < time.time())),
        'active_sessions': session_stats,
        'metrics': api_metrics
    })

async def metrics_handler(request: web.Request) -> web.Response:
    """Handler for detailed metrics endpoint (protected)"""
    # Check for API key in header or query param
    auth_header = request.headers.get('Authorization', '')
    api_key = None
    
    if auth_header.startswith('Bearer '):
        api_key = auth_header[7:]
    else:
        api_key = request.query.get('key')
    
    # Validate API key (using SECRET_TOKEN as the API key)
    if not api_key or api_key != SECRET_TOKEN:
        return web.json_response({
            'error': 'Unauthorized'
        }, status=401)
    
    # Get detailed metrics
    detailed_metrics = metrics_tracker.get_detailed_metrics()
    
    return web.json_response(detailed_metrics)

async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
    # Check if maintenance mode is enabled
    if MAINTENANCE_MODE:
        # Return an error response indicating maintenance mode
        return web.json_response({
            'error': 'Server is in maintenance mode',
            'maintenance': True
        }, status=503)  # 503 Service Unavailable
    
    ws = web.WebSocketResponse(
        max_msg_size=1024*1024*20,  # 20MB max message size
        timeout=30.0  # we want to keep things tight and short
    )
    
    await ws.prepare(request)
    
    # Get the Hugging Face token from query parameters
    hf_token = request.query.get('hf_token', '')
    
    # Generate a unique user ID for this connection
    user_id = str(uuid.uuid4())
    
    # Validate the token and determine the user role
    user_role = await session_manager.shared_api.validate_user_token(hf_token)
    logger.info(f"User {user_id} connected with role: {user_role}")
    
    # Get client IP address
    peername = request.transport.get_extra_info('peername')
    if peername is not None:
        client_ip = peername[0]
    else:
        client_ip = request.headers.get('X-Forwarded-For', 'unknown').split(',')[0].strip()
    
    logger.info(f"Client {user_id} connecting from IP: {client_ip} with role: {user_role}")
    
    # Check for anonymous user connection limits
    if user_role == 'anon':
        async with anon_connection_lock:
            # Track this connection
            anon_connections[client_ip] = anon_connections.get(client_ip, 0) + 1
            # Store the IP so we can clean up later
            ws.client_ip = client_ip
            
            # Log multiple connections from same IP but don't restrict them
            if anon_connections[client_ip] > 1:
                logger.info(f"Multiple anonymous connections from IP {client_ip}: {anon_connections[client_ip]} connections")
    
    # Store the user role in the websocket for easy access
    ws.user_role = user_role
    ws.user_id = user_id
    
    # Register with metrics
    metrics_tracker.register_session(user_id, client_ip)
    
    # Create a new session for this user
    user_session = await session_manager.create_session(user_id, user_role, ws)

    try:
        async for msg in ws:
            if msg.type == WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                    action = data.get('action')
                    
                    # Check for rate limiting
                    request_type = 'other'
                    if action in ['join_chat', 'leave_chat', 'chat_message']:
                        request_type = 'chat'
                    elif action in ['generate_video']:
                        request_type = 'video'
                    elif action == 'search':
                        request_type = 'search'
                    elif action == 'simulate':
                        request_type = 'simulation'
                    
                    # Record the request for metrics
                    await metrics_tracker.record_request(user_id, client_ip, request_type, user_role)
                    
                    # Check rate limits (except for admins)
                    if user_role != 'admin' and await metrics_tracker.is_rate_limited(user_id, request_type, user_role):
                        await ws.send_json({
                            'action': action,
                            'requestId': data.get('requestId'),
                            'success': False,
                            'error': f'Rate limit exceeded for {request_type} requests. Please try again later.'
                        })
                        continue
                    
                    # Route requests to appropriate queues
                    if action in ['join_chat', 'leave_chat', 'chat_message']:
                        await user_session.chat_queue.put(data)
                    elif action in ['generate_video']:
                        request_id = data.get('requestId', 'unknown')
                        #logger.info(f"[{request_id}] Received generate_video request from user {user_id}, adding to video queue")
                        await user_session.video_queue.put(data)
                    elif action == 'search':
                        await user_session.search_queue.put(data)
                    elif action == 'simulate':
                        await user_session.simulation_queue.put(data)
                    else:
                        await user_session.process_generic_request(data)
                        
                except Exception as e:
                    logger.error(f"Error processing WebSocket message for user {user_id}: {str(e)}")
                    await ws.send_json({
                        'action': data.get('action') if 'data' in locals() else 'unknown',
                        'success': False,
                        'error': f'Error processing message: {str(e)}'
                    })
                    
            elif msg.type in (WSMsgType.ERROR, WSMsgType.CLOSE):
                break
                
    finally:
        # Cleanup session
        await session_manager.delete_session(user_id)
        
        # Cleanup anonymous connection tracking
        if getattr(ws, 'user_role', None) == 'anon' and hasattr(ws, 'client_ip'):
            client_ip = ws.client_ip
            async with anon_connection_lock:
                if client_ip in anon_connections:
                    anon_connections[client_ip] = max(0, anon_connections[client_ip] - 1)
                    if anon_connections[client_ip] == 0:
                        del anon_connections[client_ip]
                    logger.info(f"Anonymous connection from {client_ip} closed. Remaining: {anon_connections.get(client_ip, 0)}")
        
        # Unregister from metrics
        metrics_tracker.unregister_session(user_id, client_ip)
        logger.info(f"Connection closed for user {user_id}")
    
    return ws

async def init_app() -> web.Application:
    app = web.Application(
        client_max_size=1024**2*20  # 20MB max size
    )
    
    # Add cleanup logic
    async def cleanup(app):
        logger.info("Shutting down server, closing all sessions...")
        await session_manager.close_all_sessions()
    
    app.on_shutdown.append(cleanup)
    
    # Add routes
    app.router.add_get('/ws', websocket_handler)
    app.router.add_get('/api/status', status_handler)
    app.router.add_get('/api/metrics', metrics_handler)
    
    # Set up static file serving
    # Define the path to the public directory
    public_path = Path(__file__).parent / 'build' / 'web'
    if not public_path.exists():
        public_path.mkdir(parents=True, exist_ok=True)
    
    # Set up static file serving with proper security considerations
    async def static_file_handler(request):
        # Get the path from the request (removing leading /)
        path_parts = request.path.lstrip('/').split('/')
        
        # Convert to safe path to prevent path traversal attacks
        safe_path = public_path.joinpath(*path_parts)
        
        # Make sure the path is within the public directory (prevent directory traversal)
        try:
            safe_path = safe_path.resolve()
            if not str(safe_path).startswith(str(public_path.resolve())):
                return web.HTTPForbidden(text="Access denied")
        except (ValueError, FileNotFoundError):
            return web.HTTPNotFound()
        
        # If path is a directory, look for index.html
        if safe_path.is_dir():
            safe_path = safe_path / 'index.html'
        
        # Check if the file exists
        if not safe_path.exists() or not safe_path.is_file():
            # If not found, serve index.html (for SPA routing)
            safe_path = public_path / 'index.html'
            if not safe_path.exists():
                return web.HTTPNotFound()
        
        # Determine content type based on file extension
        content_type = 'text/plain'
        ext = safe_path.suffix.lower()
        if ext == '.html':
            content_type = 'text/html'
        elif ext == '.js':
            content_type = 'application/javascript'
        elif ext == '.css':
            content_type = 'text/css'
        elif ext in ('.jpg', '.jpeg'):
            content_type = 'image/jpeg'
        elif ext == '.png':
            content_type = 'image/png'
        elif ext == '.gif':
            content_type = 'image/gif'
        elif ext == '.svg':
            content_type = 'image/svg+xml'
        elif ext == '.json':
            content_type = 'application/json'
        
        # Return the file with appropriate headers
        return web.FileResponse(safe_path, headers={'Content-Type': content_type})
    
    # Add catch-all route for static files (lower priority than API routes)
    app.router.add_get('/{path:.*}', static_file_handler)
    
    return app

if __name__ == '__main__':
    app = asyncio.run(init_app())
    web.run_app(app, host='0.0.0.0', port=8080)