main.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
from contextlib import asynccontextmanager

from fastapi import FastAPI, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from granian import Granian

from app.admin import api as admin_api
from app.admin import routes as admin_routes
from app.core import claude, openai
from app.core.config import settings
from app.core.upstream import UpstreamClient
from app.utils.logger import setup_logger
from app.utils.reload_config import RELOAD_CONFIG

# Setup logger
logger = setup_logger(log_dir="logs", debug_mode=settings.DEBUG_LOGGING)


async def warmup_upstream_client():
    """Optional warmup for upstream adapters to initialize dynamic dependencies early."""
    try:
        client = UpstreamClient()
        logger.info(
            f"✅ Upstream adapter ready, supported models: {len(client.get_supported_models())}"
        )
    except Exception as exc:
        logger.warning(f"⚠️ Upstream adapter warmup failed: {exc}")


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Initialize Token Database
    from app.services.request_log_dao import init_request_log_dao
    from app.services.token_automation import (
        run_directory_import,
        start_token_automation_scheduler,
        stop_token_automation_scheduler,
    )
    from app.services.token_dao import init_token_database

    await init_token_database()
    init_request_log_dao()

    if (
        settings.TOKEN_AUTO_IMPORT_ENABLED
        and settings.TOKEN_AUTO_IMPORT_SOURCE_DIR.strip()
    ):
        try:
            await run_directory_import(
                settings.TOKEN_AUTO_IMPORT_SOURCE_DIR,
                provider="zai",
            )
            logger.info("✅ Startup directory auto-import completed")
        except Exception as exc:
            logger.warning(f"⚠️ Startup directory auto-import failed: {exc}")

    # Initialize authentication token pool from database
    from app.utils.token_pool import initialize_token_pool_from_db

    token_pool = await initialize_token_pool_from_db(
        provider="zai",
        failure_threshold=settings.TOKEN_FAILURE_THRESHOLD,
        recovery_timeout=settings.TOKEN_RECOVERY_TIMEOUT,
    )

    if not token_pool and not settings.ANONYMOUS_MODE:
        logger.warning(
            "⚠️ No available tokens found and anonymous mode is disabled, the service may not work correctly"
        )

    if settings.ANONYMOUS_MODE:
        from app.utils.guest_session_pool import initialize_guest_session_pool

        guest_pool = await initialize_guest_session_pool(
            pool_size=settings.GUEST_POOL_SIZE,
        )
        guest_status = guest_pool.get_pool_status()
        logger.info(
            "🫥 Anonymous session pool ready: "
            f"{guest_status.get('valid_sessions', 0)} available sessions"
        )

    await warmup_upstream_client()
    await start_token_automation_scheduler()

    yield

    logger.info("🔄 Application shutting down...")

    await stop_token_automation_scheduler()

    if settings.ANONYMOUS_MODE:
        from app.utils.guest_session_pool import close_guest_session_pool

        await close_guest_session_pool()


# Create FastAPI app with lifespan
# root_path is used for reverse proxy path prefix (e.g., /api or /path-prefix)
app = FastAPI(lifespan=lifespan, root_path=settings.ROOT_PATH)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
    allow_headers=["Content-Type", "Authorization"],
)

# Mount web static files directory
try:
    app.mount("/static", StaticFiles(directory="app/static"), name="static")
except RuntimeError:
    # If static directory doesn't exist, create it
    os.makedirs("app/static/css", exist_ok=True)
    os.makedirs("app/static/js", exist_ok=True)
    app.mount("/static", StaticFiles(directory="app/static"), name="static")

# Include API routers
app.include_router(openai.router)
app.include_router(claude.router)

# Include admin routers
app.include_router(admin_routes.router)
app.include_router(admin_api.router)


@app.options("/")
async def handle_options():
    """Handle OPTIONS requests"""
    return Response(status_code=200)


@app.get("/")
async def root():
    """Root endpoint"""
    return {"message": "OpenAI Compatible API Server"}


def run_server():
    service_name = settings.SERVICE_NAME

    logger.info(f"🚀 Starting {service_name} service...")
    logger.info(f"📡 Listening on: 0.0.0.0:{settings.LISTEN_PORT}")
    logger.info(f"🔧 Debug Mode: {'Enabled' if settings.DEBUG_LOGGING else 'Disabled'}")
    logger.info(f"🔐 Anonymous Mode: {'Enabled' if settings.ANONYMOUS_MODE else 'Disabled'}")

    try:
        Granian(
            "main:app",
            interface="asgi",
            address="0.0.0.0",
            port=settings.LISTEN_PORT,
            reload=False,  # Disable hot reload in production
            process_name=service_name,  # Set process name
            **RELOAD_CONFIG,  # Hot reload configuration
        ).serve()
    except KeyboardInterrupt:
        logger.info("🛑 Received interrupt signal, shutting down service...")
    except Exception as e:
        logger.error(f"❌ Service startup failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    run_server()