from __future__ import annotations

from app.config import get_settings
from app.db import get_connection
from app.embeddings import embed_query
from app.llm import synthesize_answer
from app.models import SearchResponse, SourceChunk


def retrieve_chunks(query_embedding: list[float], k: int) -> list[SourceChunk]:
    sql = """
        SELECT
            id,
            text_chunk,
            source_file,
            chunk_index,
            1 - (embedding <=> %s::vector) AS similarity
        FROM document_chunks
        ORDER BY embedding <=> %s::vector
        LIMIT %s
    """
    with get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, (query_embedding, query_embedding, k))
            rows = cur.fetchall()
    return [
        SourceChunk(
            id=row["id"],
            text_chunk=row["text_chunk"],
            similarity=float(row["similarity"]),
            chunk_index=row["chunk_index"],
            source_file=row["source_file"],
        )
        for row in rows
    ]


def answer_question(question: str) -> SearchResponse:
    settings = get_settings()
    embedding = embed_query(question)
    sources = retrieve_chunks(embedding, settings.retrieval_k)
    if not sources:
        return SearchResponse(
            answer="This is not covered in the provided document.",
            sources=[],
        )
    answer = synthesize_answer(question, [s.text_chunk for s in sources])
    return SearchResponse(answer=answer, sources=sources)