"""Candidate selection agent for the playground's retrieval path.

Given a question and a wide, source-diverse pool of retrieved chunks, an agent
decides which chunks contain the information the question needs. Rank order
alone cannot do this: on the evaluation corpus, the answering chunk sat as
deep as fused rank 63 while near-duplicates and on-vocabulary neighbors filled
the top slots. The agent reads the pool and selects; it does not answer.

Candidates are sent in full — no truncation. The 90-candidate pool measured
about 32k tokens on the evaluation corpus, well inside the model context, and
any cut risks hiding the one row of a large table that answers the question.
Labels are local indices that carry no information about any expected answer.
"""

from __future__ import annotations

import asyncio
from functools import lru_cache
from typing import TYPE_CHECKING, List

from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.settings import ModelSettings

from quber.agents.langsmith_tracer import usage_metadata_from
from quber.playground.agent import anthropic_model
from quber.playground.tracing import run_metadata, tracer

if TYPE_CHECKING:
    from quber.playground.retrieval import RetrievedChunk

MODEL = "claude-haiku-4-5"

SYSTEM_PROMPT = """\
You are a retrieval triage step for financial-document question answering.
You are given a QUESTION and a numbered list of CANDIDATE context chunks
retrieved from one document. Your ONLY job is to decide which candidates
contain the specific information needed to answer the question.

Do NOT answer the question. Do NOT compute or state any figure. Select
candidates.

Rules:
- Return the labels (e.g. C0, C7) of every candidate that contains information
  needed to answer the question, ordered MOST relevant first.
- A candidate is relevant if it holds the exact line, row, or statement the
  question asks about (matching the entity, period, metric, and units). Prefer
  the candidate that states the asked-for value most directly and specifically.
- If several near-duplicate candidates exist, include the one(s) that most
  precisely match the question's entity and period; you may include a few close
  alternates but keep the best first.
- Only return labels that appear in the candidate list. Never invent a label.
- If nothing is relevant, return an empty list.
"""


class Selection(BaseModel):
    labels: List[str] = Field(
        default_factory=list,
        description="Candidate labels (e.g. 'C3') that answer the question, most relevant first.",
    )


@lru_cache(maxsize=1)
def _agent() -> Agent[None, Selection]:
    # Built over whichever credential the host presents, the same way the
    # answer model is. The hosted task carries only the API key; demanding the
    # OAuth token here meant selection never ran there and every hosted answer
    # came from fused order alone, with nothing but a log line to say so.
    return Agent(
        anthropic_model(MODEL),
        output_type=Selection,
        system_prompt=SYSTEM_PROMPT,
        model_settings=ModelSettings(temperature=0.0),
    )


def _prompt(question: str, cands: List["RetrievedChunk"]) -> str:
    blocks = [
        f"[label=C{i} | type {c.chunk_type} | page {c.page + 1}]\n{c.content}" for i, c in enumerate(cands)
    ]
    return (
        f"QUESTION: {question}\n\n"
        f"CANDIDATES ({len(cands)}):\n\n" + "\n\n".join(blocks) + "\n\n"
        "Return the labels of the candidates that contain the information needed "
        "to answer the question, most relevant first."
    )


async def select(question: str, cands: List["RetrievedChunk"]) -> List[int]:
    """Indices into `cands` the agent judged relevant, most relevant first.

    Labels outside the candidate range and duplicates are discarded, so the
    result is always a valid, de-duplicated index list.
    """
    prompt = _prompt(question, cands)
    inputs = {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ],
        "chunk_ids": [c.chunk_id for c in cands],
    }
    async with tracer().llm_run("select_chunks", inputs, model=MODEL, extra_metadata=run_metadata()) as run:
        result = await _agent().run(prompt)
        run.outputs = {
            "messages": [{"role": "assistant", "content": result.output.model_dump_json()}],
            "usage_metadata": usage_metadata_from(result.usage),
        }
    idxs: List[int] = []
    for lab in result.output.labels:
        lab = lab.strip().lstrip("Cc")
        if lab.isdigit():
            j = int(lab)
            if 0 <= j < len(cands) and j not in idxs:
                idxs.append(j)
    return idxs


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