"""Query endpoint for natural language Q&A."""

import asyncio
import time
from uuid import UUID

import structlog
from fastapi import APIRouter, HTTPException
from pymilvus import connections
from sqlalchemy import select

from src.api.dependencies import ApiKeyDep, AsyncSessionDep
from src.api.tenant_context import TenantContextDep, set_rls_tenant
from src.api.schemas.query import QueryRequest, QueryResponse
from src.core.exceptions import DocumentNotFoundError, QueryError, RetrievalError
from src.db.models import Document
from src.retrieval.answer_generator import AnswerGenerator
from src.retrieval.semantic_retriever import SemanticRetriever

logger = structlog.get_logger(__name__)

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


async def validate_document_ids(
    document_ids: list[UUID],
    session: AsyncSessionDep,
    tenant_id: UUID | None = None,
    node_id: UUID | None = None,
) -> None:
    """Validate that all document IDs exist and are completed.

    Args:
        document_ids: List of document IDs to validate
        session: Database session
        tenant_id: Tenant UUID for scoped validation
        node_id: Node UUID for scoped validation

    Raises:
        DocumentNotFoundError: If any document doesn't exist or isn't completed
    """
    for doc_id in document_ids:
        query = select(Document).where(
            Document.id == doc_id,
            Document.status == "completed",
        )
        if tenant_id is not None:
            query = query.where(Document.tenant_id == tenant_id)
        if node_id is not None:
            query = query.where(Document.node_id == node_id)
        result = await session.execute(query)
        doc = result.scalar_one_or_none()
        if not doc:
            raise DocumentNotFoundError(
                f"Document {doc_id} not found or not processed",
                details={"document_id": str(doc_id)},
            )


async def process_query(
    request: QueryRequest,
    session: AsyncSessionDep,
    tenant_id: str = "",
    node_id: str = "",
) -> QueryResponse:
    """Process a query request.

    Args:
        request: Query request with query text and optional filters
        session: Database session

    Returns:
        Query response with answer, sources, and metadata

    Raises:
        DocumentNotFoundError: If document_ids are invalid
        RetrievalError: If retrieval fails
        QueryError: If answer generation fails
    """
    start_time = time.time()

    # Validate document IDs if provided (scoped by tenant + node)
    if request.document_ids:
        from uuid import UUID as _UUID
        _tid = _UUID(tenant_id) if tenant_id else None
        _nid = _UUID(node_id) if node_id else None
        await validate_document_ids(request.document_ids, session, _tid, _nid)

    # Retrieve relevant sources
    retriever = SemanticRetriever(top_k=5)

    # Convert tag_filters to dict if provided
    tag_filter_dict = None
    if request.tag_filters:
        tag_filter_dict = {}
        if request.tag_filters.contains_any:
            tag_filter_dict["contains_any"] = request.tag_filters.contains_any
        if request.tag_filters.contains_all:
            tag_filter_dict["contains_all"] = request.tag_filters.contains_all

    sources, top_similarity = await retriever.retrieve(
        query=request.query,
        tenant_id=tenant_id,
        node_id=node_id,
        document_ids=request.document_ids,
        tag_filters=tag_filter_dict,
    )

    # Generate answer
    generator = AnswerGenerator()
    answer, confidence = await generator.generate(
        query=request.query,
        sources=sources if request.include_context else [],
        top_similarity=top_similarity,
    )

    # Calculate processing time
    processing_time_ms = int((time.time() - start_time) * 1000)

    return QueryResponse(
        query=request.query,
        answer=answer,
        confidence=confidence,
        sources=sources if request.include_context else [],
        processing_time_ms=processing_time_ms,
    )


@router.post(
    "",
    status_code=200,
    response_model=QueryResponse,
    responses={
        200: {
            "description": "Query processed successfully",
            "model": QueryResponse,
        },
        400: {
            "description": "Invalid request",
            "content": {
                "application/json": {
                    "examples": {
                        "empty_query": {
                            "summary": "Empty query",
                            "value": {
                                "detail": [
                                    {
                                        "loc": ["body", "query"],
                                        "msg": "Query cannot be empty",
                                        "type": "value_error",
                                    }
                                ]
                            },
                        },
                        "query_too_long": {
                            "summary": "Query exceeds length limit",
                            "value": {
                                "detail": [
                                    {
                                        "loc": ["body", "query"],
                                        "msg": "ensure this value has at most 1000 characters",
                                        "type": "value_error.any_str.max_length",
                                    }
                                ]
                            },
                        },
                    }
                }
            },
        },
        401: {
            "description": "Authentication required",
        },
        404: {
            "description": "Document not found",
            "content": {
                "application/json": {
                    "example": {
                        "error": {
                            "code": "NOT_FOUND",
                            "message": "Document {id} not found or not processed",
                            "details": {"document_id": "uuid"},
                        }
                    }
                }
            },
        },
        504: {
            "description": "Query processing timed out",
            "content": {
                "application/json": {
                    "example": {
                        "error": {
                            "code": "TIMEOUT",
                            "message": "Query processing timed out (>5 seconds)",
                        }
                    }
                }
            },
        },
    },
)
async def query_documents(
    request: QueryRequest,
    api_key: ApiKeyDep,
    session: AsyncSessionDep,
    ctx: TenantContextDep,
) -> QueryResponse:
    """Process a natural language query against uploaded documents.

    Query documents using semantic search and return AI-generated answers
    with source references from the knowledge base.

    **Features:**
    - Semantic search using vector similarity
    - Natural language answer generation
    - Source references with document context
    - Optional document filtering
    - Processing time tracking

    **Example Request:**
    ```json
    {
      "query": "What does item code 2024-4_0019 appear in?",
      "document_ids": ["uuid1", "uuid2"],
      "include_context": true
    }
    ```

    **Example Response:**
    ```json
    {
      "query": "What does item code 2024-4_0019 appear in?",
      "answer": "Item code 2024-4_0019 appears in the Screen column...",
      "confidence": 0.92,
      "sources": [
        {
          "file": "distribution_spec.xlsx",
          "sheet": "Item Mapping",
          "location": "A42:D42",
          "context": "2024-4_0019 | - | ○ | Screen display item"
        }
      ],
      "processing_time_ms": 1250
    }
    ```
    """
    try:
        logger.info(
            "query_received",
            query=request.query,
            document_ids=request.document_ids,
            include_context=request.include_context,
            api_key_id=str(api_key.id),
        )

        # Set RLS session variable for PostgreSQL defense-in-depth
        await set_rls_tenant(session, ctx.tenant_id)

        # Process query with 5-second timeout
        try:
            response = await asyncio.wait_for(
                process_query(
                    request, session,
                    tenant_id=str(ctx.tenant_id),
                    node_id=str(ctx.node_id),
                ),
                timeout=5.0,
            )

            logger.info(
                "query_completed",
                query=request.query,
                processing_time_ms=response.processing_time_ms,
                confidence=response.confidence,
                sources_count=len(response.sources),
            )

            return response

        except asyncio.TimeoutError:
            logger.warning("query_timeout", query=request.query)
            raise HTTPException(
                status_code=504,
                detail={
                    "error": {
                        "code": "TIMEOUT",
                        "message": "Query processing timed out (>5 seconds)",
                    }
                },
            )

    except DocumentNotFoundError as e:
        logger.warning(
            "document_not_found",
            query=request.query,
            error=str(e),
            details=e.details,
        )
        raise HTTPException(status_code=404, detail=e.to_dict()["error"])

    except (RetrievalError, QueryError) as e:
        logger.error(
            "query_processing_error",
            query=request.query,
            error=str(e),
            error_code=e.code,
        )
        raise HTTPException(status_code=500, detail=e.to_dict()["error"])

    except Exception as e:
        logger.exception(
            "query_unexpected_error",
            query=request.query,
            error=str(e),
        )
        # Ensure Milvus connection is cleaned up
        try:
            connections.disconnect("default")
        except Exception:
            pass
        raise HTTPException(
            status_code=500,
            detail={
                "error": {
                    "code": "INTERNAL_ERROR",
                    "message": "An unexpected error occurred",
                }
            },
        )
