"""Answering into a fixed envelope whose payload is a closed set of shapes.

The agent always returns `Answer`. It chooses the payload — scalar, series,
grid, prose, unanswerable — from the union, so deciding whether the question
wants a figure or an explanation happens inside the same call that answers it.
No separate classification step, and no schema is invented at runtime: every
shape a caller can receive is written down in `shapes`.
"""

from __future__ import annotations

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

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.answers.shapes import Answer
from quber.playground.retrieval import RetrievedChunk
from quber.playground.tracing import run_metadata, tracer

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

Answer ONLY from the provided context. Never use outside knowledge, never
invent figures. If the context does not hold the answer, return the
`unanswerable` payload.

CHOOSING THE PAYLOAD

Read what the question asks for and return that shape:

- `scalar` — the question names ONE figure. This is the common case for
  questions phrased like "Return the Q1 2026 value" or "What was X as of
  <date>". Put the figure in `value` exactly as the source prints it, with
  its currency sign and separators, and NOTHING else. "$14.47" — not
  "$14.47 per share", not "Book value was $14.47", not "$14.47 as of
  March 31, 2026". The qualifiers belong in `unit`, `period`, and `label`.
- `series` — the question asks for one measure across several periods,
  segments, or categories.
- `grid` — the question asks for a whole table.
- `prose` — the question asks you to explain, compare, summarize, or
  reason. Only here may you write sentences.
- `unanswerable` — the context does not contain what was asked.

When a question could be read either way, prefer the narrower shape: if a
single figure answers it, return `scalar`.

GROUNDING

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 — for tighter grounding — the finer-grained element ids exposed inside a
chunk's content, such as table cells tagged `<td id="t0-8-1">`. On a `scalar`,
also set `source_id` to the single cell or chunk the figure was read from, and
make `value` character-for-character what that cell prints.

Only cite ids that appear verbatim in the context. Never invent or construct
an id.
"""


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


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"
        "Answer from the context above, choosing the payload shape the question asks for."
    )
    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
) -> Answer:
    prompt = prompt_for(question, chunks, document)
    async with tracer().llm_run(
        "answer_fixed", trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata()
    ) as run:
        result = await build_agent().run(prompt)
        output: Answer = result.output
        run.outputs = {
            "messages": [{"role": "assistant", "content": output.model_dump_json()}],
            "usage_metadata": usage_metadata_from(result.usage),
        }
    return output