"""pick_for_question — one constrained LLM call that returns a typed Pick.

Used by L2 Analyses that need to translate a free-form question into the
composer's input. The model is shown:
  - the narrowed metric catalog (name + description + unit)
  - the narrowed candidate column refs (bare or `table.column`)
  - the filter grammar (via the system prompt)

It returns JSON. We validate the JSON into a `Pick` and fail fast on any
shape error — no local retry. Per the project's retry feedback, recovery
belongs at a future global layer that emits a visible event when it fires.

If the model returns `{"error": "..."}` (the prompt's escape hatch for
"this question can't be expressed"), we raise PickValidationError with
the model's reason — the caller surfaces it cleanly.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass

from api import langfuse_client as lf
from api.analyses._narrow import Candidates, narrow_candidates
from api.composer.types import Pick, PickValidationError
from api.llm import chat
from api.prompts import load, render
from api.recon import load_recon
from api.recon.types import Recon


@dataclass
class PickOutcome:
    pick: Pick
    candidates: Candidates  # the narrowed set the model saw (kept for tracing)


def pick_for_question(
    question: str,
    *,
    recon: Recon | None = None,
    allowed_metrics: list[str] | None = None,
    max_tokens: int = 512,
    span_name: str = "pick.gen",
) -> PickOutcome:
    """Run one LLM call to translate `question` into a Pick.

    `allowed_metrics`: optional list to restrict the candidate metrics
    (used by Analyses that already know which metric the planner chose,
    so we don't waste tokens listing every metric in the catalog).
    """
    recon = recon or load_recon()
    candidates = narrow_candidates(recon, allowed_metrics=allowed_metrics)

    metrics_block = _render_metrics_block(candidates)
    columns_block = _render_columns_block(candidates)

    system = load("pick.system")
    user = render(
        "pick.user",
        question=question,
        metrics_block=metrics_block,
        columns_block=columns_block,
    )

    with lf.span(
        "pick_for_question",
        input={
            "question": question,
            "metric_candidates": [m.name for m in candidates.metrics],
            "column_candidate_count": len(candidates.column_refs),
        },
    ) as span:
        raw = chat(system=system, user=user, max_tokens=max_tokens, span_name=span_name)
        payload = _parse_json(raw)
        if "error" in payload:
            raise PickValidationError(
                f"LLM declined to pick: {payload['error']!s}"
            )
        pick = Pick.from_dict(payload)
        # Bind-check the pick against recon up-front — the composer would
        # raise the same errors later, but raising here makes the failure
        # event happen at the pick step instead of compose.
        _validate_against_recon(pick, recon)
        span.update(output=pick.to_dict())
        return PickOutcome(pick=pick, candidates=candidates)


def _render_metrics_block(candidates: Candidates) -> str:
    if not candidates.metrics:
        return "(no candidate metrics)"
    lines: list[str] = []
    for m in candidates.metrics:
        unit = f" [{m.unit}]" if m.unit else ""
        lines.append(f"- {m.name}{unit}: {m.description}")
    return "\n".join(lines)


def _render_columns_block(candidates: Candidates) -> str:
    if not candidates.column_refs:
        return "(no candidate columns)"
    return "\n".join(f"- {c}" for c in candidates.column_refs)


def _parse_json(text: str) -> dict:
    """Best-effort: prefer a ```json``` fence, otherwise extract the first
    {...} block. Matches the existing pattern in drill_down._parse_json."""
    m = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
    raw = m.group(1) if m else text
    start, end = raw.find("{"), raw.rfind("}")
    if start < 0 or end <= start:
        raise PickValidationError(
            f"pick_for_question returned no JSON object: {text[:200]!r}"
        )
    try:
        return json.loads(raw[start:end + 1])
    except json.JSONDecodeError as e:
        raise PickValidationError(
            f"pick_for_question returned invalid JSON: {e!s}; raw={raw[:200]!r}"
        ) from e


def _validate_against_recon(pick: Pick, recon: Recon) -> None:
    """Surface metric/column resolution errors as PickValidationError up
    front so the trace marks the pick step (not compose) as the failure."""
    if pick.metric not in recon.metrics:
        raise PickValidationError(
            f"picked metric {pick.metric!r} is not in recon"
        )
    for ref in pick.group_by:
        try:
            recon.resolve_column(ref)
        except ValueError as e:
            raise PickValidationError(
                f"group_by column {ref!r}: {e}"
            ) from e
    for f in pick.where:
        try:
            recon.resolve_column(f.column)
        except ValueError as e:
            raise PickValidationError(
                f"where filter on {f.column!r}: {e}"
            ) from e