"""Streaming a typed answer.

The playground streams today because a prose answer is worth watching arrive.
A figure is not — by the time a scalar validates there is nothing to animate —
but the question is not whether streaming is useful for a scalar, it is
whether adding a shape breaks the streaming path that prose still needs.

So the stream yields the same two events as before. `delta` carries whatever
text is readable so far, which is the prose payload's text while it grows and
nothing at all for a scalar; `final` carries the validated `Answer`. A caller
that only wants the value ignores the deltas and waits for the final event.
"""

from __future__ import annotations

from typing import AsyncIterator, List, Optional, Tuple

from quber.agents.langsmith_tracer import usage_metadata_from
from quber.playground.agent import DEFAULT_MODEL
from quber.playground.answers.document import DocumentIdentity
from quber.playground.answers.fixed import build_agent, prompt_for, trace_inputs
from quber.playground.answers.shapes import Answer
from quber.playground.retrieval import RetrievedChunk
from quber.playground.tracing import run_metadata, tracer


def readable(partial: object) -> str:
    """The text worth showing from a partially built answer, if any.

    Only the prose payload has text that grows. A partial whose payload is not
    yet determined, or is a figure, has nothing to show.
    """
    payload = getattr(partial, "payload", None)
    if payload is None:
        return ""
    if getattr(payload, "kind", None) != "prose":
        return ""
    return getattr(payload, "text", "") or ""


async def answer_stream(
    question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None
) -> AsyncIterator[Tuple[str, object]]:
    """Yield ("delta", text) as prose accumulates, then ("final", 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:
        async with build_agent().run_stream(prompt) as result:
            last = ""
            async for partial in result.stream_output(debounce_by=0.05):
                text = readable(partial)
                if text and text != last:
                    last = text
                    yield "delta", text
            final: Answer = await result.get_output()
            run.outputs = {
                "messages": [{"role": "assistant", "content": final.model_dump_json()}],
                "usage_metadata": usage_metadata_from(result.usage),
            }
        yield "final", final