import asyncio
import logging
import os
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from qdrant_client import AsyncQdrantClient

from aggregation import aggregate_message_ids
from config import (
    API_KEY,
    HOST,
    PORT,
    QDRANT_URL,
    validate_required_env,
    logger,
)
from query_builder import (
    build_primary_query,
    build_extra_dense_queries,
    build_sparse_query,
    embed_dense,
    embed_dense_batch,
    embed_sparse,
    get_sparse_model,
)
from rerank import rerank_points
from retrieval import qdrant_search
from schemas import SearchAPIItem, SearchAPIRequest, SearchAPIResponse

validate_required_env()


@asynccontextmanager
async def lifespan(app: FastAPI):
    await asyncio.to_thread(get_sparse_model)
    logger.info("BM25 model preloaded")
    app.state.http = httpx.AsyncClient(
        timeout=30.0,
        limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
    )
    app.state.qdrant = AsyncQdrantClient(url=QDRANT_URL, api_key=API_KEY)
    try:
        yield
    finally:
        await app.state.http.aclose()
        await app.state.qdrant.close()


app = FastAPI(
    title="Search Service",
    version="0.1.0",
    lifespan=lifespan,
)


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


@app.post("/search", response_model=SearchAPIResponse)
async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
    question = payload.question
    query = question.text.strip()
    if not query:
        raise HTTPException(status_code=400, detail="question.text is required")

    client: httpx.AsyncClient = app.state.http
    qdrant: AsyncQdrantClient = app.state.qdrant

    primary_query = build_primary_query(question)
    sparse_query = build_sparse_query(question)
    extra_queries = build_extra_dense_queries(question)

    dense_task = embed_dense(client, primary_query)
    sparse_task = asyncio.to_thread(embed_sparse, sparse_query)
    primary_dense, sparse_vector = await asyncio.gather(dense_task, sparse_task)

    extra_dense: list[list[float]] = []
    if extra_queries:
        try:
            extra_dense = await embed_dense_batch(client, extra_queries[:2])
        except Exception as e:
            logger.warning("Extra dense embedding failed: %s", e)

    all_points = await qdrant_search(qdrant, primary_dense, extra_dense, sparse_vector, question)

    if not all_points:
        return SearchAPIResponse(results=[])

    reranked_head, tail = await rerank_points(client, query, all_points)
    message_ids = aggregate_message_ids(reranked_head, tail)

    return SearchAPIResponse(results=[SearchAPIItem(message_ids=message_ids)])


@app.exception_handler(Exception)
async def exception_handler(request: Request, exc: Exception) -> JSONResponse:
    logger.exception(exc)
    detail = str(exc) or repr(exc)

    if isinstance(exc, RequestValidationError):
        return JSONResponse(status_code=422, content={"detail": exc.errors()})

    if isinstance(exc, HTTPException):
        return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})

    return JSONResponse(status_code=500, content={"detail": detail})


def main() -> None:
    import uvicorn

    uvicorn.run("main:app", host=HOST, port=PORT, reload=False)


if __name__ == "__main__":
    main()