from __future__ import annotations

import logging
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING

import sentry_sdk
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from sentry_sdk.integrations.fastapi import FastApiIntegration

from codeflash_api._config import settings
from codeflash_api.db._engine import create_pool
from codeflash_api.db._queries import Queries
from codeflash_api.llm._client import LLMClient
from codeflash_api.observability._background import drain

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

log = logging.getLogger(__name__)


@asynccontextmanager
async def _lifespan(app: FastAPI) -> AsyncIterator[dict[str, object]]:
    log.info("Starting codeflash-api (%s)", settings.environment)

    pool = await create_pool(settings.database_url)
    queries = Queries(pool=pool)

    app.state.queries = queries
    app.state.rate_limit_cache = {}
    app.state.llm_client = LLMClient()

    state: dict[str, object] = {}
    yield state

    await drain()
    await pool.close()
    log.info("Shutting down codeflash-api")


def create_app() -> FastAPI:
    app = FastAPI(
        title="Codeflash API",
        version="0.1.0",
        lifespan=_lifespan,
        docs_url="/docs" if settings.debug else None,
        redoc_url=None,
    )

    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )

    if settings.sentry_dsn:
        sentry_sdk.init(
            dsn=settings.sentry_dsn,
            integrations=[FastApiIntegration()],
            environment=settings.environment,
        )

    _register_routes(app)

    return app


def _register_routes(app: FastAPI) -> None:
    from codeflash_api.adaptive._router import (
        router as adaptive_router,
    )
    from codeflash_api.explain._router import (
        router as explain_router,
    )
    from codeflash_api.jit._router import (
        router as jit_router,
    )
    from codeflash_api.logging._router import (
        router as logging_router,
    )
    from codeflash_api.optimize._router import (
        router as optimize_router,
    )
    from codeflash_api.ranking._router import (
        router as ranking_router,
    )
    from codeflash_api.refinement._router import (
        router as refinement_router,
    )
    from codeflash_api.repair._router import (
        router as repair_router,
    )
    from codeflash_api.review._router import (
        router as review_router,
    )
    from codeflash_api.testgen._review_router import (
        router as testgen_review_router,
    )
    from codeflash_api.testgen._router import (
        router as testgen_router,
    )
    from codeflash_api.workflow._router import (
        router as workflow_router,
    )

    app.include_router(optimize_router)
    app.include_router(adaptive_router)
    app.include_router(explain_router)
    app.include_router(testgen_router)
    app.include_router(testgen_review_router)
    app.include_router(jit_router)
    app.include_router(ranking_router)
    app.include_router(refinement_router)
    app.include_router(logging_router)
    app.include_router(repair_router)
    app.include_router(review_router)
    app.include_router(workflow_router)

    @app.get("/healthcheck")
    async def healthcheck() -> dict[str, str]:
        return {"status": "ok"}
