"""Plan layer — one LLM call that selects Analyses for a question.

Public surface is framework-free. Returns a `Plan` that the runtime iterates.
"""

from __future__ import annotations

import json
import logging
import re
from typing import Any

from api import langfuse_client as lf
from api.analyses.registry import catalog
from api.llm import chat
from api.plan.types import Plan
from api.prompts import load, render
from api.tools.schema import load_schema

logger = logging.getLogger("nvi.plan.planner")


def plan(question: str) -> Plan:
    schema = load_schema()
    user = render(
        "planner.user",
        question=question,
        table_names=", ".join(schema.table_names()),
        metrics_block=schema.render_metrics(),
        catalog=json.dumps(catalog(), indent=2),
    )

    with lf.span("plan", as_type="generation", input={"question": question}) as span:
        result = Plan.from_dict(
            _parse(chat(system=load("planner.system"), user=user, max_tokens=1024))
        )
        span.update(output={
            "rationale": result.rationale,
            "steps": [s.to_dict() for s in result.steps],
        })
        if not result.steps:
            raise ValueError("planner returned empty plan")
        return result


def _parse(text: str) -> dict[str, Any]:
    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 ValueError(f"planner returned no JSON: {text[:200]!r}")
    return json.loads(raw[start:end + 1])