"""LLM-driven NL → SQL with sqlglot validation and one retry."""

from __future__ import annotations

import logging
import re

import sqlglot

from api import langfuse_client as lf
from api.llm import chat
from api.prompts import load, render
from api.recon import load_recon
from api.recon.validate import ReconValidationError, validate_sql
from api.tools.types import T2SResult

logger = logging.getLogger("nvi.tools.text_to_sql")


def text_to_sql(question: str, *, hint_tables: list[str] | None = None) -> T2SResult:
    recon = load_recon()
    tables = hint_tables or recon.table_names()
    system = load("text_to_sql.system")

    def _user(retry_hint: str = "") -> str:
        return render(
            "text_to_sql.user",
            question=question,
            schema_block=recon.render_tables(tables),
            metrics_block=recon.render_metrics(),
            retry_hint=retry_hint,
        )

    with lf.span(
        "text_to_sql",
        as_type="generation",
        input={"question": question, "hint_tables": hint_tables},
    ) as span:
        raw = _extract_sql(chat(system=system, user=_user()))
        sql = _normalize(raw)              # raises on parse error — fail fast
        validate_sql(sql, recon)           # raises ReconValidationError — fail fast

        result = T2SResult(sql=sql, used_tables=_extract_tables(sql))
        span.update(output={"sql": sql, "used_tables": result.used_tables})
        return result


def _extract_sql(text: str) -> str:
    m = re.search(r"```sql\s*(.*?)```", text, re.DOTALL | re.IGNORECASE)
    if m:
        return m.group(1).strip().rstrip(";").strip()
    return text.strip().rstrip(";").strip()


def _normalize(sql: str) -> str:
    """Parse the LLM's SQL, then re-render with every identifier quoted.

    Postgres folds unquoted identifiers to lowercase, so `d.A13` becomes
    `d.a13` and breaks against case-preserving columns like district."A13".
    Re-rendering with `identify=True` puts double quotes around every
    identifier — case is preserved, and Postgres treats a quoted lowercase
    identifier the same as the unquoted form, so this is safe in both
    directions.
    """
    parsed = sqlglot.parse_one(sql, dialect="postgres")
    if parsed is None:
        raise ValueError("sqlglot returned no parse tree")
    return parsed.sql(dialect="postgres", identify=True)


def _extract_tables(sql: str) -> list[str]:
    try:
        parsed = sqlglot.parse_one(sql, dialect="postgres")
        return sorted({t.name for t in parsed.find_all(sqlglot.exp.Table)})
    except Exception:
        return []