"""pgvector cosine retrieval → parent-document mapping (PLAN §3/§6, A7).

Retrieval returns the *chunks* closest to the question embedding (top-K by
cosine distance). The product requirement is that the LLM receives the
**entire relevant document**, not just the chunk (LOCKED A7) — so
:meth:`select_documents` maps chunk hits back to their parent documents
(``chunks.document_id → documents``), dedupes, ranks by best chunk score,
and caps the combined context at ``BOR_MAX_CONTEXT_CHARS``.
"""
from __future__ import annotations

import uuid
from collections.abc import Sequence
from dataclasses import dataclass

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.config import get_settings
from app.models import Chunk, Document

#: Marker appended when the context budget is exceeded (PLAN §6).
TRUNCATION_MARKER = "[…truncated…]"


@dataclass
class RetrievedChunk:
    """One chunk hit: its cosine score plus the parent document row."""

    chunk_id: uuid.UUID
    position: int
    content: str
    score: float  # 1 − cosine_distance (higher is more similar)
    document: Document


def retrieve(
    db: Session, question_embedding: list[float], top_k: int | None = None
) -> list[RetrievedChunk]:
    """Top-*top_k* chunks by pgvector cosine distance (``<=>``).

    ``score = 1 − distance``. Results are ordered by ascending distance, so
    index 0 is the best hit. Chunks whose embedding is still NULL (two-phase
    import in progress) are skipped.
    """
    k = top_k if top_k is not None else get_settings().top_k_chunks
    distance = Chunk.embedding.cosine_distance(question_embedding)
    rows = db.execute(
        select(Chunk, distance.label("distance"), Document)
        .join(Document, Chunk.document_id == Document.id)
        .where(Chunk.embedding.is_not(None))
        .order_by(distance)
        .limit(k)
    ).all()
    return [
        RetrievedChunk(
            chunk_id=chunk.id,
            position=chunk.position,
            content=chunk.content,
            score=round(1.0 - float(dist), 6),
            document=doc,
        )
        for chunk, dist, doc in rows
    ]


def weak_hit_titles(chunks: Sequence[RetrievedChunk]) -> list[str]:
    """Distinct parent-document titles of *chunks*, best chunk score first.

    Deflection mode (PLAN §6, A8) is built from these *titles only* — the
    LOW prompt and the "Maybe try" chips never see document content.
    """
    titles: list[str] = []
    seen: set[uuid.UUID] = set()
    for rc in sorted(chunks, key=lambda c: c.score, reverse=True):
        if rc.document.id in seen:
            continue
        seen.add(rc.document.id)
        titles.append(rc.document.title)
    return titles


def select_documents(
    chunks: Sequence[RetrievedChunk],
    n: int | None = None,
    max_chars: int | None = None,
) -> list[Document]:
    """Map chunk hits to distinct parent documents, ranked by best chunk score.

    At most *n* documents are returned (default ``BOR_TOP_N_DOCS``). The
    returned rows carry the full document content; if the combined content
    would exceed *max_chars* (default ``BOR_MAX_CONTEXT_CHARS``), the
    lowest-ranked overflowing document is truncated in place with the
    ``[…truncated…]`` marker so the assembled context never exceeds the
    budget (PLAN §6).
    """
    top_n = n if n is not None else get_settings().top_n_docs
    budget = max_chars if max_chars is not None else get_settings().max_context_chars

    docs: list[Document] = []
    seen: set[uuid.UUID] = set()
    for rc in sorted(chunks, key=lambda c: c.score, reverse=True):
        if rc.document.id in seen:
            continue
        seen.add(rc.document.id)
        docs.append(rc.document)
    docs = docs[:top_n]

    remaining = budget
    for doc in docs:
        if len(doc.content) <= remaining:
            remaining -= len(doc.content)
        else:
            keep = max(0, remaining - len(TRUNCATION_MARKER))
            doc.content = doc.content[:keep] + TRUNCATION_MARKER
            remaining = 0
    return docs