"""SQL validation against the recon schema.

`validate_sql` parses the SQL, runs sqlglot's `qualify` optimizer pass with
the dataset's column→type schema, and raises with a clear message when the
LLM references a column that doesn't exist on the table it's bound to.

This is the deterministic complement to the LLM prompt: prompt hints help
the LLM get it right; the validator GUARANTEES we don't ship a schema-
wrong query to Postgres. If validation fails, the caller can either bubble
the error up (no silent retry) or do one *guided* re-prompt that includes
the validator's message — which is a correction with concrete facts, not
a blind retry.
"""

from __future__ import annotations

import sqlglot
from sqlglot.errors import OptimizeError
from sqlglot.optimizer.qualify import qualify

from api.recon import load_recon
from api.recon.types import Recon


class ReconValidationError(ValueError):
    """The SQL references a table/column combination that doesn't exist
    in the recon. The original sqlglot error is in `__cause__`."""


def _build_sqlglot_schema(recon: Recon) -> dict[str, dict[str, str]]:
    """Convert recon → sqlglot's expected schema shape: {table: {col: type}}."""
    return {
        name: {col.name: col.sql_type for col in t.columns}
        for name, t in recon.tables.items()
    }


def validate_sql(sql: str, recon: Recon | None = None) -> None:
    """Raise `ReconValidationError` if any column reference in `sql` doesn't
    exist on the table it's bound to. Returns None on success."""
    recon = recon or load_recon()
    schema = _build_sqlglot_schema(recon)
    try:
        parsed = sqlglot.parse_one(sql, dialect="postgres")
        qualify(parsed, schema=schema, dialect="postgres")
    except OptimizeError as e:
        # Try to enrich the message with hints from the recon.
        msg = str(e)
        hint = _column_hint(msg, recon)
        full = f"{msg}{hint}" if hint else msg
        raise ReconValidationError(full) from e


_COL_PATTERNS = [
    r"Column '([^']+)'",                       # "Column 'X' could not be resolved."
    r"Unknown column:\s*\"?([^\"\s,]+)\"?",  # "Unknown column: X"
    r"column \"([^\"]+)\"",                   # 'column "X" does not exist'
]


def _column_hint(msg: str, recon: Recon) -> str:
    """If the error names a specific column, append a hint about which
    table(s) actually own it. Tries the wordings sqlglot and psycopg use."""
    import re
    col_name: str | None = None
    for pat in _COL_PATTERNS:
        m = re.search(pat, msg)
        if m:
            col_name = m.group(1)
            break
    if col_name is None:
        return ""
    owners = recon.owning_tables(col_name)
    if owners:
        return f"  (Column '{col_name}' actually lives in: {', '.join(owners)}.)"
    return f"  (No table in the {recon.schema!r} schema has a column named '{col_name}'.)"