from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from axon.middleware.auth import APIKeyMiddleware
from axon.middleware.rate_limit import RateLimitMiddleware
from axon.middleware.request_context import RequestContextMiddleware
from axon.models import MetricsResponse
from axon.routers import crawl, extract, index, jobs, map, research, search, tavily, validate
from axon.services.cache import RedisCache, get_cache
from axon.services.metrics import metrics
from axon.services.reranker import get_reranker, load_reranker
from axon.services.index import init_index


@asynccontextmanager
async def lifespan(app: FastAPI):
    await init_index()
    load_reranker()
    yield


app = FastAPI(title="Axon Search API", version="1.0.0", lifespan=lifespan)

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

app.include_router(search.router)
app.include_router(extract.router)
app.include_router(crawl.router)
app.include_router(map.router)
app.include_router(index.router)
app.include_router(jobs.router)
app.include_router(research.router)
app.include_router(tavily.router)
app.include_router(validate.router)


@app.get("/health")
async def health() -> dict[str, str]:
    cache = await get_cache()
    redis_status = (
        "connected"
        if isinstance(cache, RedisCache) and await cache.ping()
        else "in-memory fallback"
    )
    reranker_status = "loaded" if get_reranker() is not None else "not loaded"
    return {
        "status": "ok",
        "redis": redis_status,
        "reranker": reranker_status,
    }


@app.get("/metrics", response_model=MetricsResponse)
async def metrics_endpoint() -> MetricsResponse:
    return MetricsResponse(**await metrics.snapshot())