core/main.py

import asyncio
import os
import traceback
from contextlib import asynccontextmanager
from importlib import import_module
from pathlib import Path

from ariadne import load_schema_from_path, make_executable_schema
from ariadne.asgi import GraphQL
from graphql import GraphQLError
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.responses import FileResponse, JSONResponse, Response
from starlette.routing import Mount, Route
from starlette.staticfiles import StaticFiles

from auth.handler import EnhancedGraphQLHTTPHandler
from auth.middleware import AuthMiddleware, auth_middleware
from auth.oauth import oauth_callback, oauth_login
from cache.precache import precache_data
from cache.revalidator import revalidation_manager
from rbac import initialize_rbac
from services.search import check_search_service, initialize_search_index, preload_models, search_service
from services.viewed import ViewedStorage
from settings import DEV_SERVER_PID_FILE_NAME
from storage.redis import redis
from storage.schema import create_all_tables, resolvers
from utils.exception import ExceptionHandlerMiddleware
from utils.logger import custom_error_formatter
from utils.logger import root_logger as logger
from utils.sentry import start_sentry

DEVMODE = os.getenv("DOKKU_APP_TYPE", "false").lower() == "false"
DIST_DIR = Path(__file__).parent / "dist"  # Директория для собранных файлов
INDEX_HTML = Path(__file__).parent / "index.html"

import_module("resolvers")

schema = make_executable_schema(load_schema_from_path("schema/"), resolvers)

# Создаем middleware с правильным порядком
middleware = [
    # Начинаем с обработки ошибок
    Middleware(ExceptionHandlerMiddleware),
    # CORS должен быть перед другими middleware для корректной обработки preflight-запросов
    Middleware(
        CORSMiddleware,
        allow_origins=[
            "https://testing.discours.io",
            "https://testing3.discours.io",
            "https://v3.dscrs.site",
            "https://session-daily.vercel.app",
            "https://coretest.discours.io",
            "https://new.discours.io",
            "https://localhost:3000",
        ],
        allow_methods=["GET", "POST", "OPTIONS"],  # Явно указываем OPTIONS
        allow_headers=["*"],
        allow_credentials=True,
    ),
    # Аутентификация должна быть после CORS
    Middleware(AuthMiddleware),
]

# Создаем экземпляр GraphQL с улучшенным обработчиком и кастомным форматтером ошибок
graphql_app = GraphQL(
    schema,
    debug=DEVMODE,
    http_handler=EnhancedGraphQLHTTPHandler(),
    error_formatter=custom_error_formatter,
)


# Оборачиваем GraphQL-обработчик для лучшей обработки ошибок


async def graphql_handler(request: Request) -> Response:
    """
    Обработчик GraphQL запросов с поддержкой middleware и обработкой ошибок.

    Выполняет:
    1. Проверку метода запроса (GET, POST, OPTIONS)
    2. Обработку GraphQL запроса через ariadne
    3. Применение middleware для корректной обработки cookie и авторизации
    4. Обработку исключений и формирование ответа

    Args:
        request: Starlette Request объект

    Returns:
        Response: объект ответа (обычно JSONResponse)
    """
    if request.method not in ["GET", "POST", "OPTIONS"]:
        return JSONResponse({"error": "Method Not Allowed by main.py"}, status_code=405)

    # Проверяем, что все необходимые middleware корректно отработали
    if not hasattr(request, "scope") or "auth" not in request.scope:
        logger.warning("[graphql] AuthMiddleware не обработал запрос перед GraphQL обработчиком")

    try:
        # Обрабатываем запрос через GraphQL приложение
        result = await graphql_app.handle_request(request)

        # Применяем middleware для установки cookie
        # Используем метод process_result из auth_middleware для корректной обработки
        # cookie на основе результатов операций login/logout
        return await auth_middleware.process_result(request, result)
    except asyncio.CancelledError:
        return JSONResponse({"error": "Request cancelled"}, status_code=499)
    except GraphQLError as e:
        # Для GraphQL ошибок (например, неавторизованный доступ) не логируем полный трейс
        logger.warning(f"GraphQL error: {e}")
        return JSONResponse({"error": str(e)}, status_code=403)
    except Exception as e:
        logger.error(f"Unexpected GraphQL error: {e!s}")
        logger.debug(f"Unexpected GraphQL error traceback: {traceback.format_exc()}")
        return JSONResponse({"error": "Internal server error"}, status_code=500)


async def spa_handler(request: Request) -> Response:
    """
    Обработчик для SPA (Single Page Application) fallback.

    Возвращает index.html для всех маршрутов, которые не найдены,
    чтобы клиентский роутер (SolidJS) мог обработать маршрутизацию.

    Args:
        request: Starlette Request объект

    Returns:
        FileResponse: ответ с содержимым index.html
    """
    # Исключаем API маршруты из SPA fallback
    path = request.url.path
    if path.startswith(("/graphql", "/oauth", "/assets")):
        return JSONResponse({"error": "Not found"}, status_code=404)

    index_path = DIST_DIR / "index.html"
    if index_path.exists():
        return FileResponse(index_path, media_type="text/html")
    return JSONResponse({"error": "Admin panel not built"}, status_code=404)


async def health_handler(request: Request) -> Response:
    """Health check endpoint with Redis monitoring"""
    try:
        redis_info = await redis.get_info()
        return JSONResponse(
            {"status": "healthy", "redis": {"connected": redis.is_connected, "ping": await redis.ping(), **redis_info}}
        )
    except Exception as e:
        logger.error(f"Health check failed: {e}")
        return JSONResponse({"status": "unhealthy", "error": str(e)}, status_code=500)


async def shutdown() -> None:
    """Остановка сервера и освобождение ресурсов"""
    logger.info("Остановка сервера")

    # Закрываем соединение с Redis
    await redis.disconnect()

    # Останавливаем поисковый сервис
    await search_service.close()

    pid_file = Path(DEV_SERVER_PID_FILE_NAME)
    if pid_file.exists():
        pid_file.unlink()


async def dev_start() -> None:
    """
    Инициализация сервера в DEV режиме.

    Функция:
    1. Проверяет наличие DEV режима
    2. Создает PID-файл для отслеживания процесса
    3. Логирует информацию о старте сервера

    Используется только при запуске сервера с флагом "dev".
    """
    try:
        pid_path = Path(DEV_SERVER_PID_FILE_NAME)
        # Если PID-файл уже существует, проверяем, не запущен ли уже сервер с этим PID
        if pid_path.exists():
            try:
                with pid_path.open(encoding="utf-8") as f:
                    old_pid = int(f.read().strip())
                # Проверяем, существует ли процесс с таким PID

                try:
                    os.kill(old_pid, 0)  # Сигнал 0 только проверяет существование процесса
                    print(f"[warning] DEV server already running with PID {old_pid}")
                except OSError:
                    print(f"[info] Stale PID file found, previous process {old_pid} not running")
            except (ValueError, FileNotFoundError):
                print("[warning] Invalid PID file found, recreating")

        # Создаем или перезаписываем PID-файл
        with pid_path.open("w", encoding="utf-8") as f:
            f.write(str(os.getpid()))
        print(f"[main] process started in DEV mode with PID {os.getpid()}")
    except Exception as e:
        logger.error(f"[main] Error during server startup: {e!s}")
        # Не прерываем запуск сервера из-за ошибки в этой функции
        print(f"[warning] Error during DEV mode initialization: {e!s}")


async def initialize_search_index_with_data() -> None:
    """Инициализация поискового индекса данными из БД"""
    try:
        from orm.shout import Shout
        from storage.db import local_session

        # Получаем все опубликованные шауты из БД
        with local_session() as session:
            shouts = session.query(Shout).filter(Shout.published_at.is_not(None)).all()

        if shouts:
            await initialize_search_index(shouts)
            print(f"[search] Loaded {len(shouts)} published shouts into search index")
        else:
            print("[search] No published shouts found to index")

    except Exception as e:
        logger.error(f"Failed to initialize search index with data: {e}")


# Глобальная переменная для background tasks
background_tasks: list[asyncio.Task] = []


@asynccontextmanager
async def lifespan(app: Starlette):
    """
    Функция жизненного цикла приложения.

    Обеспечивает:
    1. Инициализацию всех необходимых сервисов и компонентов
    2. Предзагрузку кеша данных
    3. Подключение к Redis и поисковому сервису
    4. Корректное завершение работы при остановке сервера

    Args:
        app: экземпляр Starlette приложения

    Yields:
        None: генератор для управления жизненным циклом
    """
    try:
        print("[lifespan] Starting application initialization")
        create_all_tables()

        # Инициализируем RBAC систему с dependency injection
        initialize_rbac()

        # Инициализируем Sentry для мониторинга ошибок
        start_sentry()

        await asyncio.gather(
            redis.connect(),
            precache_data(),
            ViewedStorage.init(),
            check_search_service(),
            revalidation_manager.start(),
        )
        if DEVMODE:
            await dev_start()
        print("[lifespan] Basic initialization complete")

        # Инициализируем поисковый индекс данными из БД
        print("[lifespan] Initializing search index with existing data...")
        await initialize_search_index_with_data()
        print("[lifespan] Search service initialized with Muvera")

        # 🚀 Предзагружаем ML модели после монтирования /dump
        print("[lifespan] Starting ML models preloading...")
        await preload_models()
        print("[lifespan] ML models preloading completed")

        yield
    finally:
        print("[lifespan] Shutting down application services")

        # Отменяем все background tasks
        for task in background_tasks:
            if not task.done():
                task.cancel()

        # Ждем завершения отмены tasks
        if background_tasks:
            await asyncio.gather(*background_tasks, return_exceptions=True)

        tasks = [redis.disconnect(), ViewedStorage.stop(), revalidation_manager.stop()]
        await asyncio.gather(*tasks, return_exceptions=True)
        print("[lifespan] Shutdown complete")


# Обновляем маршрут в Starlette
app = Starlette(
    routes=[
        Route("/graphql", graphql_handler, methods=["GET", "POST", "OPTIONS"]),
        # OAuth маршруты
        Route("/oauth/{provider}", oauth_login, methods=["GET"]),
        Route("/oauth/{provider}/callback", oauth_callback, methods=["GET"]),
        # Health check endpoint
        Route("/health", health_handler, methods=["GET"]),
        # Статические файлы (CSS, JS, изображения)
        Mount("/assets", app=StaticFiles(directory=str(DIST_DIR / "assets"))),
        # Корневой маршрут для админ-панели
        Route("/", spa_handler, methods=["GET"]),
        # SPA fallback для всех остальных маршрутов
        Route("/{path:path}", spa_handler, methods=["GET"]),
    ],
    middleware=middleware,  # Используем единый список middleware
    lifespan=lifespan,
    debug=True,
)

if DEVMODE:
    # Для DEV режима регистрируем дополнительный CORS middleware только для localhost
    app.add_middleware(
        CORSMiddleware,
        allow_origins=[
            "https://localhost:3000",
            "https://localhost:3001",
            "https://localhost:3002",
            "http://localhost:3000",
            "http://localhost:3001",
            "http://localhost:3002",
        ],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )