"""Answering to a stated expectation rather than an inferred one.

The union in `shapes` lets the model decide whether a question wants a figure
or an explanation, and it decides well — but "well" is not "predictably", and a
caller that wanted prose and got a figure has no recourse except to rephrase.
So the expectation is something the caller states:

    auto   the model chooses the payload from the union
    value  a figure, whatever the question looked like
    text   an explanation, whatever the question looked like

Every path returns the same `Answer`, so stating an expectation changes which
shape comes back and nothing else about the contract. `value` and `text` go
through `declared` — the caller-supplied-schema path — which is exactly what
stating an expectation means.

A forced shape can still fail honestly: asking for a figure where the context
holds none returns `unanswerable` rather than inventing one.
"""

from __future__ import annotations

from typing import AsyncIterator, List, Literal, Optional, Tuple

from quber.playground.answers import declared, fixed, streaming
from quber.playground.answers.document import DocumentIdentity
from quber.playground.answers.shapes import Answer, Prose, Scalar, Unanswerable
from quber.playground.retrieval import RetrievedChunk

Want = Literal["auto", "value", "text"]


def scalar_answer(result: declared.DeclaredScalar) -> Answer:
    if result.not_found:
        return Answer(payload=Unanswerable(reason=result.not_found), cited_ids=result.cited_ids)
    return Answer(
        payload=Scalar(value=result.value, unit=result.unit, period=result.period),
        cited_ids=result.cited_ids,
    )


def prose_answer(result: declared.DeclaredProse) -> Answer:
    if result.not_found and not result.text:
        return Answer(payload=Unanswerable(reason=result.not_found), cited_ids=result.cited_ids)
    return Answer(payload=Prose(text=result.text), cited_ids=result.cited_ids)


async def answer(
    question: str,
    chunks: List[RetrievedChunk],
    want: Want = "auto",
    document: Optional[DocumentIdentity] = None,
) -> Answer:
    if want == "value":
        return scalar_answer(await declared.answer(question, chunks, declared.DeclaredScalar, document))
    if want == "text":
        return prose_answer(await declared.answer(question, chunks, declared.DeclaredProse, document))
    return await fixed.answer(question, chunks, document)


async def answer_stream(
    question: str,
    chunks: List[RetrievedChunk],
    want: Want = "auto",
    document: Optional[DocumentIdentity] = None,
) -> AsyncIterator[Tuple[str, object]]:
    """Yield ("delta", text) while readable text grows, then ("final", Answer)."""
    if want == "text":
        async for kind, payload in declared.answer_stream(
            question, chunks, declared.DeclaredProse, "text", document
        ):
            if kind == "delta":
                yield kind, payload
            else:
                assert isinstance(payload, declared.DeclaredProse)
                yield "final", prose_answer(payload)
        return
    if want == "value":
        # Nothing to animate in a figure; one final event keeps the caller's
        # event contract identical across the three expectations.
        yield "final", await answer(question, chunks, "value", document)
        return
    async for kind, payload in streaming.answer_stream(question, chunks, document):
        yield kind, payload