"""Hybrid retrieval over the pgvector playground schema.

Two stages. First, two rankings over one document's chunks are fused:
pgvector cosine distance (meaning — a paraphrased question still lands) and
Postgres full-text search over the same content (letter — a question that
names a statement or line item verbatim finds it even when the whole-chunk
embedding dilutes the match). Reciprocal rank fusion combines them without
tuned weights; a query whose keywords match nothing degrades to the pure
vector ranking.

Second, a ranker orders the fused list, because rank order alone was
measured insufficient: on the evaluation corpus the answering chunk sat as
deep as fused rank 63 (a yield buried under same-vocabulary neighbors) or
rank 32 (a guidance row losing to on-vocabulary sibling tables), and no
scoring change recovered them. Which ranker runs is the playground's `ranker`
setting, and both satisfy `ranking.ChunkRanker`.

With `haiku`, the default, a selection agent reads a wide, source-diverse
window of the fused ranking as one prompt and picks which chunks the question
actually needs; the final context is composed from its picks, backfilled from
fused order. The agent stage answered all eleven evaluation probes,
deterministically, where fused order alone answered eight. Before the agent
sees the window, line records are capped at two per parent table — one
table's near-identical rows otherwise fill the window and crowd out other
sources; each table's whole-table chunk is never capped, so the full table
remains available even when its rows are.

With `jev`, every page and every chunk of the whole fused list is judged in
its own request, so neither the window nor the cap applies: both exist only
because an LLM reads the pool as one prompt. The context is the chunks that
cleared the ranker's cut, highest score first, then the chunks whose grading
calls failed twice, in fused order, up to `k`. Nothing else is backfilled,
because a chunk added from fused order is one Jev graded as not bearing on
the question or one on a page Jev rejected.
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from typing import Dict, List, Optional

from loguru import logger

from quber.playground import db
from quber.playground.embedding import embed_query
from quber.playground.ranking import ChunkRanker
from quber.settings import get_settings


@dataclass
class RetrievedChunk:
    chunk_id: str
    chunk_type: str
    page: int
    content: str
    score: float  # fused RRF score (0 to 2/21, about 0.095, at RRF_K = 20); higher is better
    parent_chunk_id: Optional[str] = None  # line item -> its whole-table record


# Fusion constants, grid-tested on the evaluation corpus (the arbor
# preferred-stock-dividends probe: keyword rank 2, vector rank 52). A deep
# pool lets a chunk far down one list still collect its credit from the
# other, and the lower RRF constant weights a top keyword rank enough to
# surface it. 60/40 left that probe at fused rank 16; 20/200 puts it at 8.
RRF_K = 20
POOL = 200  # candidates taken from each ranking before fusion

# The selection agent's window into the fused ranking. 90 is the measured
# minimum that contains every evaluation answer (the deepest sat at fused
# rank 63); at 60 that chunk is unreachable by any downstream step.
WINDOW = 90
# Line records kept per parent table inside the window. Without a cap, one
# large table's rows dominate the window (an 89-row table filled it almost
# alone) and the agent over-selects siblings; two per table with the
# whole-table chunk uncapped scored strictly better than no cap.
PARENT_CAP = 2
# The whole fused list, for the ranker that judges each chunk alone. It is
# the union ceiling of the two POOL-sized rankings, so no chunk either
# ranking returned is left out.
FUSED_LIMIT = 2 * POOL


async def retrieve(
    doc_key: str, query: str, k: int = 10, types: Optional[List[str]] = None
) -> List[RetrievedChunk]:
    """The playground's retrieval: the fused list, ordered by the configured ranker.

    Falls back to fused order if the ranker fails, so a provider outage
    degrades ranking quality instead of breaking retrieval.
    """
    name = get_settings().playground.ranker
    ranker: ChunkRanker
    if name == "jev":
        from quber.playground.ranking import rank

        ranker = rank
        pool = await _fused_window(doc_key, query, FUSED_LIMIT, types)
        backfill = False
    else:
        from quber.playground.selection import select

        ranker = select
        pool = _cap_line_records(await _fused_window(doc_key, query, WINDOW, types), PARENT_CAP)
        backfill = True

    try:
        picked = await ranker(query, pool)
    except Exception as exc:
        logger.warning("{} ranker failed ({}); falling back to fused order", name, exc)
        picked, backfill = [], True

    out: List[RetrievedChunk] = []
    seen: set[str] = set()
    for j in picked:
        c = pool[j]
        if c.chunk_id not in seen:
            seen.add(c.chunk_id)
            out.append(c)
        if len(out) >= k:
            break
    if backfill:
        for c in pool:
            if len(out) >= k:
                break
            if c.chunk_id not in seen:
                seen.add(c.chunk_id)
                out.append(c)
    # The fused positions say at a glance how deep the ranker reached.
    position = {c.chunk_id: n + 1 for n, c in enumerate(pool)}
    logger.debug(
        "{} ranker sent {} of {} candidates; fused positions {}",
        name,
        len(out),
        len(pool),
        [position[c.chunk_id] for c in out],
    )
    return out


def retrieve_sync(
    doc_key: str, query: str, k: int = 10, types: Optional[List[str]] = None
) -> List[RetrievedChunk]:
    return asyncio.run(retrieve(doc_key, query, k, types))


async def _fused_window(
    doc_key: str, query: str, window: int, types: Optional[List[str]] = None
) -> List[RetrievedChunk]:
    """The top `window` chunks of the fused vector+keyword ranking.

    Nothing here holds the event loop. The query embedding is a model forward
    pass on this host, compute with nothing to await, so it runs in a worker
    thread; the fused query is awaited on an async connection.
    """
    qvec = await asyncio.to_thread(embed_query, query)
    type_filter = "AND c.chunk_type = ANY(%(types)s)" if types else ""

    sql = f"""
        WITH vec AS (
            SELECT c.id, row_number() OVER (ORDER BY c.embedding <=> %(qvec)s::vector) AS rank
            FROM ade_playground.chunks c
            JOIN ade_playground.documents d ON d.id = c.document_id
            WHERE d.doc_key = %(doc_key)s AND c.embedding IS NOT NULL {type_filter}
            ORDER BY c.embedding <=> %(qvec)s::vector
            LIMIT %(pool)s
        ),
        -- OR-joined lexemes of the question: AND semantics (websearch_to_tsquery)
        -- fails whole questions, since one filler word absent from a chunk
        -- ("find", "tell") kills the match. ts_rank_cd with log-length
        -- normalization (flag 1) then rewards the chunk covering the most
        -- question terms, densest first — unnormalized, a boilerplate table
        -- repeating one query word hundreds of times (a subsidiary list
        -- repeating the company name) outranks a short exact-coverage line.
        oq AS (
            SELECT to_tsquery('english',
                COALESCE(NULLIF(array_to_string(
                    tsvector_to_array(to_tsvector('english', %(q)s)), ' | '), ''),
                'zzznomatchzzz')) AS q
        ),
        kw AS (
            SELECT c.id, row_number() OVER (
                       ORDER BY ts_rank_cd(c.content_tsv, oq.q, 1) DESC
                   ) AS rank
            FROM ade_playground.chunks c
            JOIN ade_playground.documents d ON d.id = c.document_id, oq
            WHERE d.doc_key = %(doc_key)s AND c.embedding IS NOT NULL {type_filter}
              AND c.content_tsv @@ oq.q
            LIMIT %(pool)s
        ),
        fused AS (
            SELECT COALESCE(vec.id, kw.id) AS id,
                   COALESCE(1.0 / (%(rrf)s + vec.rank), 0) +
                   COALESCE(1.0 / (%(rrf)s + kw.rank), 0) AS score
            FROM vec FULL OUTER JOIN kw ON kw.id = vec.id
        )
        SELECT c.chunk_id, c.chunk_type, c.page, c.content, fused.score, c.parent_chunk_id
        FROM fused JOIN ade_playground.chunks c ON c.id = fused.id
        ORDER BY fused.score DESC
        LIMIT %(cand)s
    """
    params = {"qvec": qvec, "doc_key": doc_key, "q": query, "pool": POOL, "rrf": RRF_K, "cand": window}
    if types:
        params["types"] = types
    async with db.connect_async() as conn:
        rows = await (await conn.execute(sql, params)).fetchall()
    return [
        RetrievedChunk(
            chunk_id=r[0],
            chunk_type=r[1],
            page=r[2],
            content=r[3],
            score=float(r[4]),
            parent_chunk_id=r[5],
        )
        for r in rows
    ]


def _cap_line_records(cands: List[RetrievedChunk], cap: int) -> List[RetrievedChunk]:
    """Keep at most `cap` line records per parent table, in fused order.

    Everything that is not a line record — prose, figures, and each table's
    whole-table chunk — passes through untouched, so a capped table's full
    content is still in the pool.
    """
    seen: Dict[str, int] = {}
    out: List[RetrievedChunk] = []
    for c in cands:
        if c.chunk_type == "line_item" and c.parent_chunk_id:
            n = seen.get(c.parent_chunk_id, 0)
            if n >= cap:
                continue
            seen[c.parent_chunk_id] = n + 1
        out.append(c)
    return out