from __future__ import annotations

from fastapi import APIRouter, Header, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse

from app.generation.llm_client import LLMClientError
from app.models.query import QueryRequest, QueryResponse
from app.models.repo_resolution import RepoClarificationRequired
from app.retrieval.retrieval_client import RetrievalClientError
from app.services.query_service import QueryService

router = APIRouter(prefix="/query", tags=["Query"])

_query_service = QueryService()


def _clarification_response(exc: RepoClarificationRequired) -> JSONResponse:
    return JSONResponse(
        status_code=409,
        content=exc.clarification.model_dump(mode="json"),
    )


@router.post("/", response_model=None)
async def query(
    body: QueryRequest,
    x_user_id: str | None = Header(default=None, alias="X-User-Id"),
    x_request_id: str | None = Header(default=None, alias="X-Request-Id"),
) -> QueryResponse | StreamingResponse | JSONResponse:
    if body.options and body.options.stream:
        try:
            await _query_service.preflight(body, request_id=x_request_id)
            return StreamingResponse(
                _query_service.stream_query(
                    body,
                    request_id=x_request_id,
                    user_id=x_user_id,
                ),
                media_type="text/event-stream",
            )
        except RepoClarificationRequired as exc:
            return _clarification_response(exc)
        except ValueError as exc:
            message = str(exc)
            if "deliverable generation is not supported with streaming" in message:
                raise HTTPException(status_code=422, detail=message) from exc
            raise HTTPException(status_code=400, detail=message) from exc

    try:
        return await _query_service.query(
            body,
            request_id=x_request_id,
            user_id=x_user_id,
        )
    except RepoClarificationRequired as exc:
        return _clarification_response(exc)
    except ValueError as exc:
        message = str(exc)
        if "question must be non-empty" in message:
            raise HTTPException(status_code=422, detail=message) from exc
        if "invalid intent_override" in message:
            raise HTTPException(status_code=422, detail=message) from exc
        if "deliverable generation is not supported with streaming" in message:
            raise HTTPException(status_code=422, detail=message) from exc
        raise HTTPException(status_code=400, detail=message) from exc
    except LookupError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc
    except RetrievalClientError as exc:
        raise HTTPException(
            status_code=503 if exc.status_code is None else exc.status_code,
            detail=str(exc),
        ) from exc
    except LLMClientError as exc:
        status = 503 if exc.status_code in (429, 504, None) else 502
        raise HTTPException(status_code=status, detail=str(exc)) from exc