import logging

from services.llm_client import LLMClient
from services.vectorstore import VectorStoreService

logger = logging.getLogger(__name__)


async def rag_query(
    vectorstore: VectorStoreService,
    llm_client: LLMClient,
    question: str,
    top_k: int = 5,
    document_ids: list[str] | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
    system_prompt: str | None = None,
) -> dict:
    """Pipeline: retrieve → augment → generate для одиночного вопроса пациента."""
    logger.info("RAG query: %s", question[:200])

    retrieved = vectorstore.query(
        query_text=question,
        top_k=top_k,
        document_ids=document_ids,
    )
    logger.info("Retrieved %d chunks", len(retrieved))

    llm_result = await llm_client.answer(
        question=question,
        sources=retrieved,
        system_prompt=system_prompt,
        temperature=temperature,
        max_tokens=max_tokens,
    )

    sources = []
    for item in retrieved:
        meta = item.get("metadata", {})
        sources.append({
            "document_id": meta.get("document_id", ""),
            "document_name": meta.get("document_name", ""),
            "chunk_text": item["text"][:500],
            "section": meta.get("section", ""),
            "page": meta.get("page_number", 0),
            "relevance_score": round(item.get("relevance_score", 0), 3),
        })

    return {
        "answer": llm_result["text"],
        "sources": sources,
        "model_used": llm_client.model,
        "assembled_prompt": llm_result["assembled_prompt"],
    }