"""SQL safety: AST allow-list (SELECT-only, no DDL/DML, no risky functions).

Uses sqlglot to parse and validate any LLM-generated SQL before execution.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

import sqlglot
from sqlglot import exp

ALLOWED_TABLES_PREFIXES: tuple[str, ...] = ("raw.", "dim.", "sem.")
DENY_FUNCTIONS = {"load_file", "sleep", "benchmark", "user", "version"}
MAX_LIMIT = 5000


@dataclass(frozen=True)
class SqlValidation:
    ok: bool
    reason: str = ""
    rewritten: str = ""


def validate_sql(sql: str, dialect: str = "postgres") -> SqlValidation:
    try:
        parsed = sqlglot.parse_one(sql, read=dialect)
    except Exception as e:  # noqa: BLE001
        return SqlValidation(False, f"parse_error: {e}")

    if not isinstance(parsed, (exp.Select, exp.Subqueryable, exp.With, exp.Union)):
        return SqlValidation(False, f"not_a_select: {type(parsed).__name__}")

    # Reject any DML/DDL nodes anywhere in tree.
    forbidden = (
        exp.Insert,
        exp.Update,
        exp.Delete,
        exp.Drop,
        exp.Create,
        exp.AlterTable,
        exp.TruncateTable,
        exp.Merge,
    )
    for node in parsed.walk():
        n = node[0] if isinstance(node, tuple) else node
        if isinstance(n, forbidden):
            return SqlValidation(False, f"forbidden_statement:{type(n).__name__}")
        if isinstance(n, exp.Anonymous) and (n.name or "").lower() in DENY_FUNCTIONS:
            return SqlValidation(False, f"forbidden_function:{n.name}")

    # Table allow-list
    for tbl in parsed.find_all(exp.Table):
        fq = ".".join(p for p in [tbl.db, tbl.name] if p).lower()
        if not any(fq.startswith(p) or fq == p.rstrip(".") for p in ALLOWED_TABLES_PREFIXES):
            # Also allow bare names that match dim/raw common tables (no schema qualifier)
            if tbl.name.lower() in {
                "extended_bva_v5", "bva", "item_mb", "open_po", "expense_mb",
                "opening_budget_fy_25_26", "opening_underspents", "training_blocked_budget",
                "funders", "projects", "categories", "vendors", "employees",
            }:
                continue
            return SqlValidation(False, f"forbidden_table:{fq or tbl.name}")

    # Force a LIMIT cap
    if isinstance(parsed, exp.Select):
        lim = parsed.args.get("limit")
        if lim is None:
            parsed.set("limit", exp.Limit(expression=exp.Literal.number(MAX_LIMIT)))
        else:
            try:
                n = int(lim.expression.name)
                if n > MAX_LIMIT:
                    parsed.set("limit", exp.Limit(expression=exp.Literal.number(MAX_LIMIT)))
            except Exception:
                pass

    return SqlValidation(True, rewritten=parsed.sql(dialect=dialect))