"""Retrieve API Routes"""
import uuid
from typing import List
from fastapi import APIRouter, Depends, Query
from sqlalchemy.ext.asyncio import AsyncSession
from normbase.db.connection import get_db
from normbase.models.schemas import (
    NormRetrieveRequest,
    NormRetrieveResponse,
    NormResponse,
)
from normbase.services.retrieval import RetrievalService

router = APIRouter(prefix="/retrieve", tags=["retrieve"])
_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)
):
    """Retrieve applicable norms for a given context"""
    
    retrieval_service = _get_retrieval_service()
    
    # Perform retrieval
    results, cache_hit = await retrieval_service.retrieve(
        db=db,
        context=request.context,
        agent_role=request.agent_role,
        field_ids=request.field_ids,
        top_k=request.top_k
    )
    
    # Extract norms and scores
    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),
        cache_hit=cache_hit,
    )
