"""Answering into a schema the caller declared.

No inference. The caller states the shape it wants alongside the question and
the agent is constrained to it — the same contract `Agent(output_type=...)`
already offers, applied to the retrieval path. This is what a src-package
caller does when it knows what it is asking for.

`answer` is also the shared constrained-answer call for the shapes benchmark:
`benchmark/run_shapes.py` compiles a model with `planned` and hands it here, so
the only difference between the two approaches is who authored the model.

The last rule of the system prompt, answering only from the portfolio,
statement, table or note the question names, exists because the context
now reaches deep into the fused ranking and carries look-alikes: the same
metric printed for another portfolio or period, a servicing-portfolio state
breakdown beside the loan portfolio's, a December column beside March's. The
ranker cannot tell those apart; it grades whether a chunk's words fit the
item, not which portfolio the chunk belongs to. This model reads the whole
context with its note titles and captions, so the choice sits here. Measured
on the ABR Q1 2026 10-Q sweep of 59 questions: three answers taken from the
servicing table when the question named the loan portfolio became correct,
one refusal became a correct answer, one answer to a question naming no
table in the filing became a refusal, the other 54 were unchanged.
"""

from __future__ import annotations

from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Type, TypeVar

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

from quber.agents.langsmith_tracer import usage_metadata_from
from quber.playground.agent import DEFAULT_MODEL, answer_model, build_context
from quber.playground.answers.document import DocumentIdentity, with_document
from quber.playground.retrieval import RetrievedChunk
from quber.playground.tracing import run_metadata, tracer

SYSTEM_PROMPT = """\
You are a financial-document analyst answering a question about ONE parsed
document, from context chunks retrieved from it.

Answer ONLY from the provided context. Never use outside knowledge and never
invent figures. If the context does not hold the answer, leave the answer
fields empty and say why in `not_found`.

Report figures exactly as the source prints them — currency sign, separators
and all — and nothing else. "$14.47", never "$14.47 per share" and never a
sentence. Qualifiers go in their own fields when the schema provides them.

Populate `cited_ids` with the ids supporting your answer. Each context chunk
begins with `[chunk id=<ID> | page <N> | type <TYPE>]`. Cite the chunk `<ID>`,
or the finer-grained element ids inside a chunk's content, such as table cells
tagged `<td id="t0-8-1">`. Only cite ids that appear verbatim in the context.
Never invent or construct an id.

When the question names a portfolio, statement, table or note, answer only
from that one. The same item printed for a different portfolio, statement or
table is not the answer: if the named one does not print the item, say so in
`not_found` and do not substitute the other.
"""

M = TypeVar("M", bound=BaseModel)


class DeclaredScalar(BaseModel):
    """The shape a caller declares when it wants one figure back.

    Carried here because it is the model the comparison hands to this path for
    a value-seeking question, and the shape a src caller would most often want.
    """

    value: str = Field(
        description=(
            "The figure exactly as printed in the source, with its currency sign "
            "and separators and nothing else: '$14.47', '4.6x', '(0.02)'."
        )
    )
    unit: Optional[str] = Field(default=None, description="What the number counts.")
    period: Optional[str] = Field(default=None, description="The period or as-of date, as printed.")
    cited_ids: List[str] = Field(
        default_factory=list, description="Chunk and/or cell ids supporting the value."
    )
    not_found: Optional[str] = Field(
        default=None, description="Set only when the context lacks the answer; say what."
    )


class DeclaredProse(BaseModel):
    """The shape a caller declares when it wants an explanation.

    Sibling to `DeclaredScalar`: the two of them are what the playground's
    expectation toggle picks between.
    """

    text: str = Field(description="The explanation, in prose, drawn only from the context.")
    cited_ids: List[str] = Field(
        default_factory=list, description="Chunk and/or cell ids supporting the explanation."
    )
    not_found: Optional[str] = Field(
        default=None, description="Set only when the context lacks the answer; say what."
    )


def prompt_for(
    question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None
) -> str:
    prompt = (
        f"Context chunks:\n\n{build_context(chunks)}\n\n"
        f"Question: {question}\n\n"
        "Fill the answer fields from the context above."
    )
    return with_document(prompt, document)


def build_agent(output_type: Type[M]) -> Agent[None, M]:
    return Agent(
        answer_model(),
        output_type=NativeOutput(output_type),
        system_prompt=SYSTEM_PROMPT,
        model_settings=ModelSettings(temperature=0.0),
    )


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],
    output_type: Type[M],
    document: Optional[DocumentIdentity] = None,
) -> M:
    prompt = prompt_for(question, chunks, document)
    async with tracer().llm_run(
        "answer_declared", trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata()
    ) as run:
        result = await build_agent(output_type).run(prompt)
        output: M = 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],
    output_type: Type[M],
    text_field: str,
    document: Optional[DocumentIdentity] = None,
) -> AsyncIterator[Tuple[str, object]]:
    """Yield ("delta", text) as one named field accumulates, then ("final", model).

    `text_field` names the field worth watching arrive. A declared scalar has
    none worth streaming, so callers wanting a figure skip this.
    """
    prompt = prompt_for(question, chunks, document)
    async with tracer().llm_run(
        "answer_declared", trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata()
    ) as run:
        async with build_agent(output_type).run_stream(prompt) as result:
            last = ""
            async for partial in result.stream_output(debounce_by=0.05):
                text = getattr(partial, text_field, "") or ""
                if text and text != last:
                    last = text
                    yield "delta", text
            final: M = await result.get_output()
            run.outputs = {
                "messages": [{"role": "assistant", "content": final.model_dump_json()}],
                "usage_metadata": usage_metadata_from(result.usage),
            }
        yield "final", final