"""Node: plan_query — emit a structured plan before SQL generation.

Two-stage decomposition (DIN-SQL / MAC-SQL style). The plan node produces a
JSON skeleton (tables/joins/filters/group/agg/projection/expected_row_count)
that the downstream `generate_sql` node sees as additional grounding context.

Empirically, forcing the model to commit to row-shape and projection BEFORE
writing SQL fixes a large fraction of "row_count_off" and "projection_diff"
failures observed in the BIRD baseline taxonomy (see scripts/error_taxonomy.py).
"""

from __future__ import annotations

from collections.abc import Callable

from nl_sql.agent.nodes._support import (
    render_fewshot_block,
    render_schema_block,
)
from nl_sql.agent.prompts import load_prompt
from nl_sql.agent.state import PipelineState
from nl_sql.llm.providers.base import GenerateRequest, LLMProvider


def make_plan_node(
    provider: LLMProvider,
    *,
    max_tokens: int = 600,
    temperature: float = 0.0,
    sort_schema_block: bool = False,
) -> Callable[[PipelineState], PipelineState]:
    def node(state: PipelineState) -> PipelineState:
        question = state.get("question", "")
        dialect = state.get("dialect", "sqlite")
        context = state.get("context")
        prompt = load_prompt(
            "plan",
            dialect=dialect,
            schema_block=render_schema_block(context, sort_alphabetically=sort_schema_block),
            fewshot_block=render_fewshot_block(context),
            question=question,
        )
        response = provider.generate(
            GenerateRequest(prompt=prompt, max_tokens=max_tokens, temperature=temperature)
        )
        plan_text = (response.text or "").strip()
        trace = list(state.get("trace") or [])
        trace.append(
            {
                "node": "plan_query",
                "model": response.model,
                "input_tokens": response.input_tokens,
                "output_tokens": response.output_tokens,
            }
        )
        return {
            "plan": plan_text,
            "trace": trace,
        }

    return node