"""Answering into a schema built for the question that was asked.

Two calls. A planner reads the question alone — never the context, never the
answer — and declares the fields the answer should have. Those declarations
are compiled into a Pydantic model, and the answering agent is constrained to
it through the same call `declared` uses. The only difference between the two
approaches is who authored the model.

The planner emits field declarations rather than raw JSON Schema. A model
writing free-form JSON Schema produces dialects the validator then rejects,
and that surfaces as a retry loop rather than an answer; a closed set of field
types cannot fail that way.
"""

from __future__ import annotations

from functools import lru_cache
from typing import Any, Dict, List, Literal, Optional, Tuple, Type

from pydantic import BaseModel, Field, create_model
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 anthropic_model
from quber.playground.tracing import run_metadata, tracer

PLANNER_MODEL = "claude-haiku-4-5"

FieldType = Literal["string", "number", "string_list", "number_list"]

_PY_TYPES: Dict[str, Any] = {
    "string": str,
    "number": float,
    "string_list": List[str],
    "number_list": List[float],
}

PLANNER_PROMPT = """\
You design the return shape for a financial-document question. You are given
the QUESTION only — never the document, never the answer. Do not attempt to
answer it.

Decide first whether the question asks for data or for explanation.

- If it asks you to explain, compare, summarize, or reason, set `wants_value`
  false and return no fields. It will be answered in prose.
- Otherwise set `wants_value` true and declare the fields that hold the
  answer, and only those.

Rules for fields:
- Declare the narrowest set that answers the question. A question naming one
  figure gets ONE field holding that figure, plus at most the qualifiers the
  question itself names.
- Name fields in snake_case after what they hold.
- Each field's description must tell the answering agent to report the figure
  exactly as printed in the source, with no surrounding words.
- Never declare a field for citations or sources; those are added for you.
- Never declare a field whose value would be a sentence.
"""


class FieldSpec(BaseModel):
    name: str = Field(description="snake_case field name.")
    type: FieldType = Field(description="The field's type.")
    description: str = Field(description="What the answering agent must put here.")
    required: bool = True


class Plan(BaseModel):
    wants_value: bool = Field(
        description="True when the question asks for data; false when it asks for explanation."
    )
    fields: List[FieldSpec] = Field(
        default_factory=list, description="The answer's fields. Empty when wants_value is false."
    )


@lru_cache(maxsize=1)
def _planner() -> Agent[None, Plan]:
    return Agent(
        anthropic_model(PLANNER_MODEL),
        output_type=NativeOutput(Plan),
        system_prompt=PLANNER_PROMPT,
        model_settings=ModelSettings(temperature=0.0),
    )


async def plan_for(question: str) -> Plan:
    prompt = f"QUESTION: {question}"
    inputs = {
        "messages": [
            {"role": "system", "content": PLANNER_PROMPT},
            {"role": "user", "content": prompt},
        ]
    }
    async with tracer().llm_run(
        "plan_answer", inputs, model=PLANNER_MODEL, extra_metadata=run_metadata()
    ) as run:
        result = await _planner().run(prompt)
        output: Plan = result.output
        run.outputs = {
            "messages": [{"role": "assistant", "content": output.model_dump_json()}],
            "usage_metadata": usage_metadata_from(result.usage),
        }
    return output


def compile_model(plan: Plan) -> Type[BaseModel]:
    """Turn a plan's field declarations into a model the answer agent fills.

    `cited_ids` and `not_found` are appended to every compiled model, so
    grounding and the not-present case survive whatever the planner declared.
    A plan that declared a field by either of those names loses it to the
    appended one, which is the intended precedence.
    """
    fields: Dict[str, Tuple[Any, Any]] = {}
    for spec in plan.fields:
        if spec.name in ("cited_ids", "not_found"):
            continue
        py = _PY_TYPES[spec.type]
        if spec.required:
            fields[spec.name] = (py, Field(description=spec.description))
        else:
            fields[spec.name] = (Optional[py], Field(default=None, description=spec.description))
    fields["cited_ids"] = (
        List[str],
        Field(default_factory=list, description="Chunk and/or cell ids supporting the answer."),
    )
    fields["not_found"] = (
        Optional[str],
        Field(default=None, description="Set only when the context lacks the answer; say what."),
    )
    return create_model("PlannedAnswer", **fields)  # type: ignore[call-overload]