"""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.tools.schema import load_schema
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:
    schema = load_schema()
    tables = hint_tables or schema.table_names()
    system = load("text_to_sql.system")

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

    with lf.span(
        "text_to_sql",
        as_type="generation",
        input={"question": question, "hint_tables": hint_tables},
    ) as span:
        sql = _extract_sql(chat(system=system, user=_user()))
        try:
            _validate(sql)
        except Exception as e:
            logger.info("t2s first attempt invalid (%s); retrying once", e)
            hint = f"\n\nPrevious attempt failed parsing: {e}. Fix and return only the SQL."
            sql = _extract_sql(chat(system=system, user=_user(hint)))
            _validate(sql)

        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 _validate(sql: str) -> None:
    parsed = sqlglot.parse_one(sql, dialect="postgres")
    if parsed is None:
        raise ValueError("sqlglot returned no parse tree")


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 []