"""The playground's Claude model factory, its answer-context format, and the prose baseline.

`anthropic_model` decides which credential a playground agent runs on. Every
Pydantic AI agent defined in this package builds its model through it: the
answering agents, the selection agent in `selection.py`, and the planner in
`answers/planned.py`. Two agents the playground also runs, the cover-page
attributor at upload and the footnote lookup at ingest, are built in
`quber.agents` and choose their credential there. `answer_model` fixes which Claude model answers, for the
live answering agents in `answers/fixed.py` and `answers/declared.py` and for
the prose agent here. `build_context` is the chunk-to-prompt format the
answering agents share.

The prose agent (`answer`, `GroundedAnswer`) answers in free text, strictly
from the retrieved chunks, and returns the ids it grounded on. A table chunk's
HTML carries an id on every cell (`<td id="t0-8-1">`), so it can cite
individual cells. It is not on the playground's request path.
`benchmark/run_shapes.py` uses `answer` as its baseline, and `answer_stream`
has no caller.
"""

from __future__ import annotations

from functools import lru_cache
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.models.anthropic import AnthropicModel
from pydantic_ai.providers.anthropic import AnthropicProvider
from pydantic_ai.settings import ModelSettings

from quber.agents._oauth_gate import make_oauth_anthropic_model
from quber.agents.langsmith_tracer import usage_metadata_from
from quber.playground.answers.document import DocumentIdentity, with_document
from quber.playground.retrieval import RetrievedChunk
from quber.playground.tracing import run_metadata, tracer
from quber.settings import get_settings

# The Claude model that answers.
DEFAULT_MODEL = "claude-opus-5"


def anthropic_model(model_id: str):
    """An Anthropic model over whichever credential the host presents.

    The subscription OAuth token wins when it is set, routed through the OAuth
    gate in `quber.agents._oauth_gate`. Without one the API key is used through the stock model,
    which is how the hosted playground runs: the token never leaves the
    developer host, and the task is given only the key. Neither set is a
    configuration error, reported as such.
    """
    llm = get_settings().llm
    if llm.anthropic_auth_token:
        return make_oauth_anthropic_model(model_id, llm.anthropic_auth_token)
    if llm.anthropic_api_key:
        return AnthropicModel(model_id, provider=AnthropicProvider(api_key=llm.anthropic_api_key))
    raise RuntimeError(
        "Neither ANTHROPIC_AUTH_TOKEN nor ANTHROPIC_API_KEY is set; the playground needs one of them."
    )


def answer_model():
    """The playground's answering model.

    One place decides which model answers. The fixed and declared answering
    agents and the prose baseline below all build from here rather than each
    pinning its own copy of the name.
    """
    return anthropic_model(DEFAULT_MODEL)


SYSTEM_PROMPT = """\
You are a financial-document analyst answering questions about ONE parsed
document. You are given a set of context chunks retrieved from that document.

Rules:
- Answer ONLY from the provided context chunks. Never use outside knowledge
  and never invent figures. If the context does not contain the answer, say
  so plainly.
- Be precise and quantitative. Report exact values, units, and periods as
  they appear in the source.
- Populate `cited_ids` with the ids that support your answer. Each context
  chunk begins with a header line: `[chunk id=<ID> | page <N> | type <TYPE>]`.
  * By default, cite the `<ID>` from the header of each chunk you used.
  * If a chunk's content exposes finer-grained element ids — e.g. table cells
    tagged like `<td id="0-8">` — you MAY cite those specific cell ids instead,
    for tighter grounding.
  * Only cite ids that literally appear in the provided context (a header
    `<ID>` or an element id inside a chunk's content). NEVER invent, guess, or
    construct an id that is not present verbatim.
"""


class GroundedAnswer(BaseModel):
    answer: str = Field(description="The answer, drawn only from the context chunks.")
    cited_ids: List[str] = Field(
        default_factory=list,
        description="Chunk ids and/or table cell ids that support the answer.",
    )


@lru_cache(maxsize=1)
def _agent() -> Agent[None, GroundedAnswer]:
    return Agent(
        answer_model(),
        output_type=GroundedAnswer,
        system_prompt=SYSTEM_PROMPT,
        model_settings=ModelSettings(temperature=0.0),
    )


def build_context(chunks: List[RetrievedChunk]) -> str:
    blocks = []
    for c in chunks:
        # Chunk pages are 0-based; show 1-based to match the PDF viewer.
        header = f"[chunk id={c.chunk_id} | page {c.page + 1} | type {c.chunk_type}]"
        blocks.append(f"{header}\n{c.content}")
    return "\n\n".join(blocks)


def _prompt(question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None) -> str:
    context = build_context(chunks)
    prompt = (
        f"Context chunks:\n\n{context}\n\n"
        f"Question: {question}\n\n"
        "Answer from the context above and cite the ids you used."
    )
    return with_document(prompt, document)


def _trace_inputs(prompt: str, chunks: List[RetrievedChunk]) -> Dict[str, Any]:
    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ],
        "chunk_ids": [c.chunk_id for c in chunks],
    }


async def answer(
    question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None
) -> GroundedAnswer:
    prompt = _prompt(question, chunks, document)
    async with tracer().llm_run(
        "answer_grounded", _trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata()
    ) as run:
        result = await _agent().run(prompt)
        output: GroundedAnswer = result.output
        run.outputs = {
            "messages": [{"role": "assistant", "content": output.model_dump_json()}],
            "usage_metadata": usage_metadata_from(result.usage),
        }
    return output


async def answer_stream(
    question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None
):
    """Yield ("delta", <answer so far>) as the model generates, then
    ("final", GroundedAnswer) once the full output validates."""
    prompt = _prompt(question, chunks, document)
    async with tracer().llm_run(
        "answer_grounded", _trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata()
    ) as run:
        async with _agent().run_stream(prompt) as result:
            last = ""
            async for partial in result.stream_output(debounce_by=0.05):
                text = getattr(partial, "answer", "") or ""
                if text and text != last:
                    last = text
                    yield "delta", text
            final: GroundedAnswer = await result.get_output()
            run.outputs = {
                "messages": [{"role": "assistant", "content": final.model_dump_json()}],
                "usage_metadata": usage_metadata_from(result.usage),
            }
        yield "final", final