"""Node: re-ask the LLM for SQL given the previous failure context.

Per arch v2 §3, exactly one repair pass per question. The graph guards this
by checking ``state["repair_attempted"]`` before routing here. This node
sets the flag itself before returning so a second routing attempt would be
a programming error in the graph.
"""

from __future__ import annotations

from collections.abc import Callable

from nl_sql.agent.nodes._support import (
    parse_generate_sql_output,
    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_repair_once_node(
    provider: LLMProvider,
    *,
    max_tokens: int = 1024,
    temperature: float = 0.0,
    sort_schema_block: bool = False,
) -> Callable[[PipelineState], PipelineState]:
    def node(state: PipelineState) -> PipelineState:
        generated = state.get("generated")
        previous_sql = generated.sql if generated else ""
        error_context = state.get("last_error") or "(no error context — likely a programming bug)"
        question = state.get("question", "")
        dialect = state.get("dialect", "sqlite")
        context = state.get("context")

        plan_raw = (state.get("plan") or "").strip()
        plan_block = plan_raw if plan_raw else "(no plan — repair SQL directly)"
        prompt = load_prompt(
            "repair_sql",
            dialect=dialect,
            schema_block=render_schema_block(context, sort_alphabetically=sort_schema_block),
            question=question,
            previous_sql=previous_sql,
            error_context=error_context,
            plan_block=plan_block,
        )
        response = provider.generate(
            GenerateRequest(prompt=prompt, max_tokens=max_tokens, temperature=temperature)
        )
        parsed = parse_generate_sql_output(response.text)

        trace = list(state.get("trace") or [])
        trace.append(
            {
                "node": "repair_once",
                "model": response.model,
                "confidence": parsed.confidence,
                "previous_error": error_context,
                "input_tokens": response.input_tokens,
                "output_tokens": response.output_tokens,
            }
        )
        return {
            "generated": parsed,
            "outcome": None,
            "repair_attempted": True,
            "last_error": "",
            "error_kind": None,
            "error_message": "",
            "trace": trace,
        }

    return node