"""Retrieval API Routes."""
from fastapi import APIRouter, Depends, Header
from sqlalchemy.ext.asyncio import AsyncSession

from normbase.db.connection import get_db
from normbase.schemas import (
    NormRetrieveRequest,
    NormRetrieveResponse,
    NormResponse,
)
from normbase.services.retrieval_service import RetrievalService

router = APIRouter(prefix="/retrieve", tags=["retrieval"])
_retrieval_service: RetrievalService | None = None


def _get_retrieval_service() -> RetrievalService:
    global _retrieval_service
    if _retrieval_service is None:
        _retrieval_service = RetrievalService()
    return _retrieval_service


def prewarm_retrieval_service() -> None:
    """Eagerly build RetrievalService to avoid first-request cold start."""
    _get_retrieval_service()


@router.post("", response_model=NormRetrieveResponse)
async def retrieve_norms(
    request: NormRetrieveRequest,
    db: AsyncSession = Depends(get_db),
    x_agent_id: str | None = Header(default=None),
    x_key_id: str | None = Header(default=None),
):
    """Retrieve applicable norms for a given context."""
    retrieval_service = _get_retrieval_service()
    results, usage_token, cache_hit, meta = await retrieval_service.retrieve_with_usage_token(
        db=db,
        context=request.context,
        query_mode=request.query_mode,
        nl_query=request.nl_query,
        agent_role=request.agent_role,
        agent_ids=request.agent_ids,
        tag_filters=request.tag_filters,
        field_ids=request.field_ids,
        top_k=request.top_k,
        agent_id=x_agent_id,
        key_id=x_key_id,
    )

    norms = []
    scores = []
    for norm, score in results:
        norms.append(NormResponse.model_validate(norm))
        scores.append(score)

    return NormRetrieveResponse(
        norms=norms,
        scores=scores,
        total=len(norms),
        usage_token=usage_token,
        cache_hit=cache_hit,
        meta=meta,
    )
