"""NL2SQL 只读 SQL 的 AST 安全校验。"""
from __future__ import annotations

from dataclasses import dataclass

from sqlglot import exp, parse
from sqlglot.errors import ParseError


_DANGEROUS_FUNCTIONS = {
    "LOAD_FILE",
    "UUID_FILE_NAME",
    "BENCHMARK",
    "SLEEP",
}

class SqlSecurityError(ValueError):
    """SQL 未通过只读和权限校验。"""


@dataclass(frozen=True)
class ValidatedSql:
    sql: str
    access_tables: set[str]


def _read_limit(statement: exp.Expression) -> int | None:
    limit = statement.args.get("limit")
    if limit is None:
        return None
    expression = limit.args.get("expression")
    if not isinstance(expression, exp.Literal) or not expression.is_number:
        raise SqlSecurityError("只允许使用数字 LIMIT")
    value = int(expression.this)
    if value < 0:
        raise SqlSecurityError("LIMIT 不能为负数")
    return value


def validate_select_sql(
    sql: str,
    *,
    authorized_tables: set[str],
    authorized_columns: dict[str, set[str]] | None = None,
    max_rows: int,
    max_joins: int = 5,
    max_columns: int = 100,
) -> ValidatedSql:
    """校验单条 SELECT，检查授权表并将 LIMIT 控制在最大行数内。"""
    if not isinstance(sql, str) or not sql.strip():
        raise SqlSecurityError("SQL 不能为空")
    if max_rows <= 0:
        raise SqlSecurityError("最大行数必须为正数")
    if max_joins < 0 or max_columns <= 0:
        raise SqlSecurityError("SQL 复杂度限制参数无效")
    try:
        statements = parse(sql, read="mysql")
    except ParseError as exc:
        raise SqlSecurityError("SQL 解析失败") from exc
    if len(statements) != 1 or not isinstance(statements[0], exp.Select):
        raise SqlSecurityError("只允许执行单条 SELECT")

    statement = statements[0]
    join_count = len(list(statement.find_all(exp.Join)))
    if join_count > max_joins:
        raise SqlSecurityError("SQL JOIN 深度超过限制")
    if len(statement.expressions) > max_columns:
        raise SqlSecurityError("SQL 返回列数超过限制")
    for function in statement.find_all(exp.Func):
        function_name = getattr(function, "name", "") or function.sql_name()
        if function_name.upper() in _DANGEROUS_FUNCTIONS:
            raise SqlSecurityError("SQL 包含危险函数")
    access_tables: set[str] = set()
    aliases: dict[str, str] = {}
    for table in statement.find_all(exp.Table):
        if table.db or table.catalog:
            raise SqlSecurityError("禁止跨库访问")
        access_tables.add(table.name)
        aliases[table.alias_or_name] = table.name
    if not access_tables.issubset(authorized_tables):
        raise SqlSecurityError("SQL 访问了未授权表")
    if authorized_columns is not None:
        if statement.find(exp.Star):
            raise SqlSecurityError("配置字段权限时禁止 SELECT *")
        default_table = next(iter(access_tables), None) if len(access_tables) == 1 else None
        for column in statement.find_all(exp.Column):
            table_name = aliases.get(column.table, column.table) or default_table
            if table_name is None or column.name not in authorized_columns.get(table_name, set()):
                raise SqlSecurityError("SQL 访问了未授权字段")

    current_limit = _read_limit(statement)
    if current_limit is None or current_limit > max_rows:
        statement = statement.limit(max_rows)
    return ValidatedSql(
        sql=statement.sql(dialect="mysql"),
        access_tables=access_tables,
    )