api/main.py

import os
import argparse
import sys
from modules.logger_tool import initialise_logger
logger = initialise_logger(__name__, os.getenv("LOG_LEVEL"), os.getenv("LOG_PATH"), 'default', True)
from fastapi import FastAPI, HTTPException
import uvicorn
import requests
from typing import Dict, Any
from modules.database.tools.neo4j_driver_tools import get_driver
from run.initialization.initialization import InitializationSystem
import time
import ssl

from run.setup import setup_cors
from run.routers import register_routes
from run.initialization import initialize_system

# FastAPI App Setup
app = FastAPI()
setup_cors(app)

# Health check endpoint
@app.get("/health")
async def health_check() -> Dict[str, Any]:
    """Health check endpoint that verifies all service dependencies"""
    health_status = {
        "status": "healthy",
        "services": {
            "neo4j": {"status": "healthy", "message": "Connected"},
            "supabase": {"status": "healthy", "message": "Connected"}
        }
    }
    
    try:
        # Check Neo4j
        driver = get_driver()
        if not driver:
            health_status["services"]["neo4j"] = {
                "status": "unhealthy",
                "message": "Failed to connect to Neo4j"
            }
            health_status["status"] = "unhealthy"
    except Exception as e:
        health_status["services"]["neo4j"] = {
            "status": "unhealthy",
            "message": f"Error checking Neo4j: {str(e)}"
        }
        health_status["status"] = "unhealthy"
    
    try:
        # Minimal check to confirm Supabase is responsive (e.g., pinging auth or storage endpoint)
        supabase_url = os.getenv("SUPABASE_URL")
        service_role_key = os.getenv("SERVICE_ROLE_KEY")
        response = requests.get(
            f"{supabase_url}/auth/v1/health",
            headers={"apikey": service_role_key},
            timeout=5
        )
        if response.status_code != 200:
            health_status["services"]["supabase"] = {
                "status": "unhealthy",
                "message": f"Supabase Auth API returned status {response.status_code}"
            }
            health_status["status"] = "unhealthy"
    except Exception as e:
        health_status["services"]["supabase"] = {
            "status": "unhealthy",
            "message": f"Error checking Supabase Auth API: {str(e)}"
        }
        health_status["status"] = "unhealthy"
    
    if health_status["status"] == "unhealthy":
        raise HTTPException(status_code=503, detail=health_status)
    
    return health_status

# Register routes
register_routes(app)

# Initialize system with retry logic
def initialize_with_retry(max_attempts: int = 3, initial_delay: int = 5) -> bool:
    """Initialize the system with retry logic"""
    attempt = 0
    delay = initial_delay
    
    while attempt < max_attempts:
        try:
            logger.info(f"Attempting system initialization (attempt {attempt + 1}/{max_attempts})")
            initialize_system()
            logger.info("System initialization completed successfully")
            return True
        except Exception as e:
            attempt += 1
            if attempt == max_attempts:
                logger.error(f"System initialization failed after {max_attempts} attempts: {str(e)}")
                return False
            
            logger.warning(f"Initialization attempt {attempt} failed: {str(e)}. Retrying in {delay} seconds...")
            time.sleep(delay)
            delay *= 2  # Exponential backoff
    
    return False

def run_initialization_mode():
    """Run only the initialization process"""
    logger.info("Running in initialization mode")
    logger.info("Starting system initialization...")
    
    if initialize_with_retry():
        logger.info("Initialization completed successfully")
        return True
    else:
        logger.error("Initialization failed after multiple attempts")
        return False

def run_development_mode():
    """Run the server in development mode with auto-reload"""
    logger.info("Running in development mode")
    logger.info("Starting uvicorn server with auto-reload...")
    
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=int(os.getenv('UVICORN_PORT', 8000)),
        log_level=os.getenv('LOG_LEVEL', 'info'),
        proxy_headers=True,
        timeout_keep_alive=10,
        reload=True
    )

def run_production_mode():
    """Run the server in production mode"""
    logger.info("Running in production mode")
    logger.info("Starting uvicorn server in production mode...")
    
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=int(os.getenv('UVICORN_PORT', 8000)),
        log_level=os.getenv('LOG_LEVEL', 'info'),
        proxy_headers=True,
        timeout_keep_alive=10,
        workers=int(os.getenv('UVICORN_WORKERS', '1'))
    )

def parse_arguments():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(
        description="ClassroomCopilot API Server",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Startup modes:
  init    - Run initialization scripts (database setup, etc.)
  dev     - Run development server with auto-reload
  prod    - Run production server (for Docker/containerized deployment)
        """
    )
    
    parser.add_argument(
        '--mode', '-m',
        choices=['init', 'dev', 'prod'],
        default='dev',
        help='Startup mode (default: dev)'
    )
    
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_arguments()
    
    # Set environment variable for backward compatibility
    if args.mode == 'dev':
        os.environ['BACKEND_DEV_MODE'] = 'true'
    else:
        os.environ['BACKEND_DEV_MODE'] = 'false'
    
    logger.info(f"Starting ClassroomCopilot API in {args.mode} mode")
    
    if args.mode == 'init':
        # Run initialization only
        success = run_initialization_mode()
        sys.exit(0 if success else 1)
    
    elif args.mode == 'dev':
        # Run development server
        run_development_mode()
    
    elif args.mode == 'prod':
        # Run production server
        run_production_mode()
    
    else:
        logger.error(f"Invalid mode: {args.mode}")
        sys.exit(1)