"""SQL validation — whitelist SELECT, blacklist dangerous operations."""

from __future__ import annotations

import re

import sqlparse
from sqlparse.sql import Statement
from sqlparse.tokens import Keyword, DML


class SQLValidationError(Exception):
    """Raised when SQL fails validation."""
    pass


# Dangerous keywords/functions that must NEVER appear in user-facing queries
_BLACKLIST = {
    "INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "TRUNCATE", "CREATE",
    "EXEC", "EXECUTE", "COPY", "GRANT", "REVOKE",
    "pg_read_file", "pg_write_file", "pg_ls_dir",
    "xp_cmdshell", "sp_configure", "sp_executesql",
    "OPENROWSET", "OPENDATASOURCE", "BULK",
    "INTO OUTFILE", "INTO DUMPFILE", "LOAD_FILE",
    "SHUTDOWN", "KILL",
}

# Patterns that indicate SQL injection attempts
_INJECTION_PATTERNS = [
    re.compile(r";\s*\w", re.IGNORECASE),           # Statement batching via semicolon
    re.compile(r"--\s*\S"),                           # Comment-based injection
    re.compile(r"/\*.*?\*/", re.DOTALL),              # Block comments
    # UNION ALL SELECT is legitimate — Claude generates these for multi-part queries.
    # Injection risk is mitigated by: single-statement check, SELECT-only check, schema whitelist.
    re.compile(r"INTO\s+OUTFILE", re.IGNORECASE),
    re.compile(r"INTO\s+DUMPFILE", re.IGNORECASE),
    re.compile(r"WAITFOR\s+DELAY", re.IGNORECASE),    # Time-based injection
    re.compile(r"BENCHMARK\s*\(", re.IGNORECASE),     # MySQL time-based
    re.compile(r"pg_sleep\s*\(", re.IGNORECASE),       # PG time-based
    re.compile(r"\\x[0-9a-fA-F]+"),                    # Hex encoded strings
]


def validate_sql(sql: str) -> str:
    """Validate and sanitize a SQL query.

    Returns the cleaned SQL if valid, raises SQLValidationError otherwise.
    """
    if not sql or not sql.strip():
        raise SQLValidationError("Empty SQL query")

    cleaned = sql.strip()

    # Remove trailing semicolons (we add our own execution boundary)
    cleaned = cleaned.rstrip(";").strip()

    # Check for injection patterns
    for pattern in _INJECTION_PATTERNS:
        if pattern.search(cleaned):
            raise SQLValidationError(
                f"SQL contains potentially dangerous pattern: {pattern.pattern}"
            )

    # Parse with sqlparse
    parsed = sqlparse.parse(cleaned)
    if not parsed:
        raise SQLValidationError("Could not parse SQL")

    if len(parsed) > 1:
        raise SQLValidationError("Multiple SQL statements are not allowed")

    stmt = parsed[0]

    # Check statement type — must be SELECT (or WITH ... SELECT)
    stmt_type = stmt.get_type()
    if stmt_type not in ("SELECT", None):
        # None can happen with CTEs (WITH ... SELECT)
        # Check if it starts with WITH
        first_token = _first_meaningful_token(stmt)
        if first_token and first_token.upper() not in ("SELECT", "WITH"):
            raise SQLValidationError(
                f"Only SELECT statements are allowed. Got: {stmt_type or first_token}"
            )

    # Check for blacklisted keywords
    sql_upper = cleaned.upper()
    for keyword in _BLACKLIST:
        # Use word boundary check to avoid false positives
        pattern = re.compile(r"\b" + re.escape(keyword) + r"\b", re.IGNORECASE)
        if pattern.search(cleaned):
            raise SQLValidationError(f"Blacklisted keyword found: {keyword}")

    # Ensure the query only touches v3q_data or v3q_results schemas
    # Block access to v3q_config (contains connection strings, etc.)
    if re.search(r"\bv3q_config\b", cleaned, re.IGNORECASE):
        raise SQLValidationError("Access to v3q_config schema is not allowed")

    # Block access to pg_catalog and information_schema from user queries
    if re.search(r"\bpg_catalog\b", cleaned, re.IGNORECASE):
        raise SQLValidationError("Access to pg_catalog is not allowed")

    if re.search(r"\binformation_schema\b", cleaned, re.IGNORECASE):
        raise SQLValidationError("Access to information_schema is not allowed")

    return cleaned


def _first_meaningful_token(stmt: Statement) -> str | None:
    """Get the first non-whitespace, non-comment token value."""
    for token in stmt.tokens:
        if token.is_whitespace:
            continue
        if token.ttype in (sqlparse.tokens.Comment.Single, sqlparse.tokens.Comment.Multiline):
            continue
        return str(token.value).split()[0] if token.value else None
    return None


def extract_tables_from_sql(sql: str) -> list[str]:
    """Extract table names referenced in a SQL query."""
    tables = set()
    # Match schema.table or just table patterns
    pattern = re.compile(
        r'(?:FROM|JOIN)\s+(?:"?(\w+)"?\.)?"?(\w+)"?',
        re.IGNORECASE,
    )
    for match in pattern.finditer(sql):
        schema = match.group(1)
        table = match.group(2)
        if table.upper() not in ("SELECT", "WHERE", "ON", "AND", "OR"):
            tables.add(f"{schema}.{table}" if schema else table)
    return sorted(tables)
