"""Jev ranking of retrieved chunks for the playground's answer path.

Ownership. This module decides which of a question's retrieved chunks the
answer model reads when the playground's ranker setting is `jev`. `selection`
holds the Haiku equivalent, and `retrieval.retrieve` chooses between them and
composes the final context. Both satisfy `ChunkRanker`: the question and the
candidates in, indices into the candidates out, most relevant first.

How it ranks. Jev, TypeSafe's scoring model, does not read the candidates as
one prompt. Every page and every chunk is judged alone, in its own request,
so the ranker can afford the whole fused list rather than a window of it.
Two passes. Each distinct page in the candidate list is asked whether it has
material relevant to the question, and a page answered no is dropped with
every candidate on it. Each candidate on a page that passed is then graded
on a four-level rubric, and `Ranking.order` keeps the ones at or above the
configured cut, highest first. Nothing is added back from fused order when
fewer clear the cut: anything added that way is a chunk Jev graded as
unrelated or same-topic, or one on a page Jev rejected.

What leaves the host. The question, each page's text and each chunk's text
are sent to TypeSafe's API, one text per request. TypeSafe states it does not
train on requests; zero data retention is an enterprise option not yet taken.

Failures. The SDK retries 429, 5xx, connection and timeout errors itself.
After each pass, the calls that still failed are run once more. A page that
fails twice counts as relevant, so its chunks are still graded. A chunk that
fails twice goes after the graded ranking in fused order. Every failed page
or chunk is named in the log and the trace, because it may be the one that
holds the answer. `rank` raises when every call in a pass fails, and on any
error that is not a `TypeSafeError`, such as a missing TYPESAFE_API_KEY or a
failed page-text query. `retrieve` then falls back to fused order.
"""

from __future__ import annotations

import asyncio
import time
from collections import Counter
from functools import lru_cache
from typing import (
    TYPE_CHECKING,
    Any,
    Awaitable,
    Callable,
    Dict,
    List,
    Optional,
    Protocol,
    TypeVar,
)

import httpx2
from loguru import logger
from pydantic import BaseModel, Field
from typesafe_sdk import (
    AsyncTypeSafeClient,
    Noul,
    NoulAnswer,
    RetryPolicy,
    Score,
    ScoreAnswer,
    TypeSafeAPIError,
    TypeSafeError,
)

from quber.agents.langsmith_tracer import usage_metadata_from
from quber.playground import db
from quber.playground.tracing import run_metadata, tracer
from quber.settings import get_settings

if TYPE_CHECKING:
    from quber.playground.retrieval import RetrievedChunk


class ChunkRanker(Protocol):
    """The question and a candidate list in; indices into that list out, most relevant first."""

    async def __call__(self, question: str, cands: List["RetrievedChunk"]) -> List[int]: ...


# The questions Jev is asked. `PAGE` passes the SDK only `instructions`, so
# the yes and no criteria are written into the question text. The wording is
# positive throughout: Jev's documentation says negated questions score worse.
PAGE_QUESTION = (
    "Does this page have material relevant to the query? "
    "Yes when the page discusses the metric, line item, or statement the query asks about. "
    "No when the page is about other subjects and nothing on it bears on the query."
)

# The chunk question is short because the levels carry the meaning. The levels
# describe situations, not degrees, and carry no numbers. Level 2 is what keeps
# a footnote that defines the metric, the components that reconcile to it, or
# the prior period column in the context the answer model reads.
CHUNK_QUESTION = "How does this chunk relate to the query?"

LEVEL_UNRELATED = "The chunk is about a different subject; nothing in it bears on the query."
LEVEL_SAME_TOPIC = (
    "The chunk is on the same topic as the query but gives neither the asked-for item "
    "nor anything that qualifies it."
)
LEVEL_SUPPORTS = (
    "The chunk supports the asked-for item: a definition, a footnote, a component, "
    "a reconciliation, or a comparison period for it."
)
LEVEL_STATES = "The chunk states the asked-for item directly, for the entity and period the query names."


# The two questions as the SDK sends them. A Noul is a yes/no question: Jev
# returns the probability of yes. A Score is a rubric: the levels are the
# criteria in order, 0 first, and Jev returns a probability per level and the
# probability-weighted level as the score. The SDK is called directly rather
# than through a PydanticAI agent because this path makes a few hundred calls
# per question, and the agent's per-call work measured three to four times
# the CPU of the call it wraps, enough to saturate one event loop under a
# batch of questions.
PAGE = Noul(instructions=PAGE_QUESTION)
CHUNK = Score(
    instructions=CHUNK_QUESTION, criteria=[LEVEL_UNRELATED, LEVEL_SAME_TOPIC, LEVEL_SUPPORTS, LEVEL_STATES]
)

# A page is relevant when Jev's probability of yes reaches this. It is where
# the answer turns from no to yes, not a tuned bar, so it is not a setting.
PAGE_THRESHOLD = 0.5


class PageScore(BaseModel):
    """One page's verdict. `page` is the stored 0-based index."""

    page: int
    relevant: bool
    # How far Jev's probability sat from the yes/no threshold, 0 to 1.
    confidence: Optional[float] = None
    error: Optional[str] = None


class ChunkScore(BaseModel):
    """One candidate's grade. `index` is its position in the candidate list."""

    index: int
    chunk_id: str
    page: int
    # The nearest level, 0 to 3.
    grade: Optional[int] = None
    # Jev's position on the rubric, the probability-weighted level. This is
    # what the cut and the ordering use: two chunks can share a grade and
    # differ in score.
    score: Optional[float] = None
    error: Optional[str] = None


class Ranking(BaseModel):
    """Everything both passes returned for one question."""

    pages: List[PageScore]
    # Every candidate on a page that passed, in candidate order.
    scores: List[ChunkScore]
    grade_cut: float
    model: Optional[str] = None
    seconds: float = 0.0
    input_tokens: int = 0
    # Failed call attempts per error class or HTTP status, after the SDK's
    # retries. Both of the ranker's attempts count, so a call that fails and
    # then succeeds adds one, and a call that fails twice adds two.
    failures: Dict[str, int] = Field(default_factory=dict)

    def order(self) -> List[int]:
        """Candidate indices at or above the cut, highest score first, then the
        chunks whose calls failed, in candidate order."""
        graded = [c for c in self.scores if c.score is not None and c.score >= self.grade_cut]
        graded.sort(key=lambda c: c.score or 0.0, reverse=True)
        failed = [c for c in self.scores if c.error is not None]
        return [c.index for c in graded] + [c.index for c in failed]

    def cleared(self) -> int:
        return sum(1 for c in self.scores if c.score is not None and c.score >= self.grade_cut)

    def kept_pages(self) -> List[int]:
        return [p.page for p in self.pages if p.relevant]


@lru_cache(maxsize=1)
def client() -> AsyncTypeSafeClient:
    """The one SDK client both passes share, so every call reuses one
    connection pool. The pool is unbounded: a batch of questions, each with
    `concurrency` calls in flight, would otherwise queue on httpx's default
    of 100 connections and reopen the ones past its keep-alive cap of 20 on
    every call. The per-question semaphore and the batch's question cap are
    what bound the load."""
    ts = get_settings().typesafe
    if not ts.api_key:
        raise RuntimeError("TYPESAFE_API_KEY is not set; the jev ranker needs it.")
    http = httpx2.AsyncClient(
        timeout=ts.timeout_seconds, limits=httpx2.Limits(max_connections=None, max_keepalive_connections=None)
    )
    return AsyncTypeSafeClient(
        api_key=ts.api_key, retry=RetryPolicy(max_retries=ts.max_retries), http_client=http
    )


async def page_texts(cands: List["RetrievedChunk"]) -> Dict[int, str]:
    """The text of each distinct candidate page: its chunks that are not line
    records, joined in stored order.

    Candidates carry no document key, so the pages are read from the
    document that holds the most of the candidates' chunk ids. Those ids are
    docling refs and positional names such as `#/texts/12`, `#/tables/3` and
    `t3-line-5`, unique only within a document. Another document that holds
    every candidate id ties with the right one, and the query does not break
    the tie, so the pages can come from the wrong document.
    """
    pages = sorted({c.page for c in cands})
    ids = [c.chunk_id for c in cands]
    sql = """
        WITH doc AS (
            SELECT document_id FROM ade_playground.chunks
            WHERE chunk_id = ANY(%(ids)s)
            GROUP BY document_id ORDER BY count(*) DESC LIMIT 1
        )
        SELECT k.page, k.content
        FROM ade_playground.chunks k JOIN doc ON doc.document_id = k.document_id
        WHERE k.page = ANY(%(pages)s) AND k.chunk_type <> 'line_item'
        ORDER BY k.page, k.id
    """
    texts: Dict[int, List[str]] = {p: [] for p in pages}
    async with db.connect_async() as conn:
        for page, content in await (await conn.execute(sql, {"ids": ids, "pages": pages})).fetchall():
            texts[page].append(content)
    return {p: "\n\n".join(parts) for p, parts in texts.items()}


T = TypeVar("T")


class _Tally:
    """What every call adds up to: tokens, the resolved model id, failures.
    Shaped so `usage_metadata_from` can read it like a run usage."""

    def __init__(self) -> None:
        self.input_tokens = 0
        self.output_tokens = 0
        self.model: Optional[str] = None
        self.failures: Counter[str] = Counter()

    def record(self, response: Any) -> None:
        self.input_tokens += response.usage.input_tokens or 0
        self.output_tokens += response.usage.output_tokens or 0
        self.model = response.model or self.model

    def failed(self, exc: Exception) -> str:
        key = str(exc.status) if isinstance(exc, TypeSafeAPIError) else type(exc).__name__
        self.failures[key] += 1
        return f"{key}: {exc}"


async def _twice(
    items: List[T], call: Callable[[T], Awaitable[Any]], failed: Callable[[Any], bool], concurrency: int
) -> List[Any]:
    """Run `call` over `items` at most `concurrency` at a time, then once more
    over the items whose result `failed`. Results in item order."""
    sem = asyncio.Semaphore(concurrency)

    async def guarded(item: T) -> Any:
        async with sem:
            return await call(item)

    results = list(await asyncio.gather(*(guarded(i) for i in items)))
    retry = [n for n, r in enumerate(results) if failed(r)]
    if retry:
        again = await asyncio.gather(*(guarded(items[n]) for n in retry))
        for n, r in zip(retry, again, strict=True):
            results[n] = r
    return results


def _levels() -> str:
    return "\n".join(
        f"{n}: {d}" for n, d in enumerate((LEVEL_UNRELATED, LEVEL_SAME_TOPIC, LEVEL_SUPPORTS, LEVEL_STATES))
    )


def _run_outputs(tally: _Tally, started: float, **fields: Any) -> Dict[str, Any]:
    """The fields every traced pass records, plus the pass's own."""
    return {
        **fields,
        "usage_metadata": usage_metadata_from(tally),
        "failures": dict(tally.failures),
        "seconds": round(time.perf_counter() - started, 2),
        "model": tally.model,
    }


async def judge_pages(question: str, cands: List["RetrievedChunk"]) -> tuple[List[PageScore], _Tally]:
    """Pass 1: one yes/no call per distinct candidate page. One traced run
    named `page_rank`, with every verdict, relevant pages first."""
    ts = get_settings().typesafe
    tally = _Tally()
    texts = await page_texts(cands)
    inputs = {
        "messages": [{"role": "system", "content": PAGE_QUESTION}, {"role": "user", "content": question}],
        "pages": [p + 1 for p in sorted(texts)],
    }

    async def judge(page: int) -> PageScore:
        text = texts.get(page, "")
        if not text:
            return PageScore(page=page, relevant=True)
        try:
            response = await client().system_one(
                f"QUERY: {question}\n\nPAGE {page + 1}:\n{text}", {"relevant": PAGE}, model=ts.model
            )
        except TypeSafeError as exc:
            return PageScore(page=page, relevant=True, error=tally.failed(exc))
        tally.record(response)
        answer = response.answers["relevant"]
        assert isinstance(answer, NoulAnswer)
        # Confidence is how far the probability sat from the threshold, 0 at
        # the threshold and 1 at either end, so a page that could go either
        # way on another call reads as near 0.
        return PageScore(
            page=page,
            relevant=answer.noul >= PAGE_THRESHOLD,
            confidence=round(abs(answer.noul - PAGE_THRESHOLD) * 2, 4),
        )

    started = time.perf_counter()
    async with tracer().llm_run("page_rank", inputs, model=ts.model, extra_metadata=run_metadata()) as run:
        pages = await _twice(sorted(texts), judge, lambda p: p.error is not None, ts.concurrency)
        judged = [p for p in pages if texts.get(p.page)]
        if judged and all(p.error is not None for p in judged):
            raise RuntimeError(f"every page call failed: {dict(tally.failures)}")
        ordered = sorted(pages, key=lambda p: (not p.relevant, p.page))
        run.outputs = _run_outputs(
            tally,
            started,
            messages=[{"role": "assistant", "content": str([p.page + 1 for p in pages if p.relevant])}],
            pages=[{**p.model_dump(), "page": p.page + 1} for p in ordered],
            kept=sum(1 for p in pages if p.relevant),
        )
    return pages, tally


async def grade_chunks(
    question: str, cands: List["RetrievedChunk"], kept: set[int]
) -> tuple[List[ChunkScore], _Tally]:
    """Pass 2: one rubric call per candidate on a kept page. One traced run
    named `chunk_rank`, with every score, best first, failed chunks last."""
    ts = get_settings().typesafe
    tally = _Tally()
    items = [(i, c) for i, c in enumerate(cands) if c.page in kept]
    inputs = {
        "messages": [
            {"role": "system", "content": f"{CHUNK_QUESTION}\n{_levels()}"},
            {"role": "user", "content": question},
        ],
        "chunk_ids": [c.chunk_id for _, c in items],
    }

    async def grade(item: tuple[int, "RetrievedChunk"]) -> ChunkScore:
        i, c = item
        prompt = f"QUERY: {question}\n\nCHUNK (page {c.page + 1}, {c.chunk_type}):\n{c.content}"
        try:
            response = await client().system_one(prompt, {"grade": CHUNK}, model=ts.model)
        except TypeSafeError as exc:
            return ChunkScore(index=i, chunk_id=c.chunk_id, page=c.page, error=tally.failed(exc))
        tally.record(response)
        answer = response.answers["grade"]
        assert isinstance(answer, ScoreAnswer)
        return ChunkScore(
            index=i,
            chunk_id=c.chunk_id,
            page=c.page,
            grade=min(int(answer.score + 0.5), len(CHUNK.criteria) - 1),
            score=answer.score,
        )

    started = time.perf_counter()
    async with tracer().llm_run("chunk_rank", inputs, model=ts.model, extra_metadata=run_metadata()) as run:
        scores = await _twice(items, grade, lambda s: s.error is not None, ts.concurrency)
        if scores and all(s.error is not None for s in scores):
            raise RuntimeError(f"every chunk call failed: {dict(tally.failures)}")
        graded = sorted((s for s in scores if s.error is None), key=lambda s: s.score or 0.0, reverse=True)
        failed = [s for s in scores if s.error is not None]
        cleared = [s for s in graded if (s.score or 0.0) >= ts.grade_cut]
        run.outputs = _run_outputs(
            tally,
            started,
            messages=[{"role": "assistant", "content": str([s.chunk_id for s in cleared])}],
            scores=[{**s.model_dump(), "page": s.page + 1} for s in graded + failed],
            grade_cut=ts.grade_cut,
            cleared=len(cleared),
            below_cut=len(graded) - len(cleared),
        )
    return scores, tally


async def score(question: str, cands: List["RetrievedChunk"]) -> Ranking:
    """Both passes over the candidates, with everything Jev returned."""
    ts = get_settings().typesafe
    started = time.perf_counter()
    pages, page_tally = await judge_pages(question, cands)
    kept = {p.page for p in pages if p.relevant}
    scores, chunk_tally = await grade_chunks(question, cands, kept)
    return Ranking(
        pages=pages,
        scores=scores,
        grade_cut=ts.grade_cut,
        model=chunk_tally.model or page_tally.model,
        seconds=time.perf_counter() - started,
        input_tokens=page_tally.input_tokens + chunk_tally.input_tokens,
        failures=dict(page_tally.failures + chunk_tally.failures),
    )


async def rank(question: str, cands: List["RetrievedChunk"]) -> List[int]:
    """Indices into `cands` in the order the answer model should read them."""
    ranking = await score(question, cands)
    order = ranking.order()
    logger.debug(
        "jev ranker: {} pages judged, {} kept; {} of {} candidates graded, {} cleared the cut of {}; "
        "{:.1f}s, {} input tokens, failures {}",
        len(ranking.pages),
        len(ranking.kept_pages()),
        len(ranking.scores),
        len(cands),
        ranking.cleared(),
        ranking.grade_cut,
        ranking.seconds,
        ranking.input_tokens,
        ranking.failures or "none",
    )
    failed_pages = [p for p in ranking.pages if p.error]
    failed_chunks = [s for s in ranking.scores if s.error]
    if failed_pages or failed_chunks:
        logger.warning(
            "jev ranker: {} page calls and {} chunk calls failed twice; pages {}; chunks {}",
            len(failed_pages),
            len(failed_chunks),
            [(p.page + 1, p.error) for p in failed_pages],
            [(s.chunk_id, s.error) for s in failed_chunks],
        )
    if ranking.failures.get("429"):
        logger.warning(
            "jev ranker: {} calls were throttled (429) past the SDK's retries; TypeSafe's limits may have changed",
            ranking.failures["429"],
        )
    return order


def rank_sync(question: str, cands: List["RetrievedChunk"]) -> List[int]:
    return asyncio.run(rank(question, cands))