"""
aaizaql.security.validator
─────────────────────────
SQLValidator: the security gate that every SQL statement must pass before
execution.  No SQL ever runs without passing through here.

Three checks (in order):
1. Prompt injection detection — scan the *question* for manipulation attempts.
2. Operation whitelist — only SELECT / WITH are allowed by default.
3. Dialect-aware structural parse — uses sqlglot to catch malformed SQL and
   detect dangerous statement types that could be disguised with comments.
"""

from __future__ import annotations

import re

import sqlglot
import structlog

from aaizaql.core.config import Settings
from aaizaql.core.exceptions import PromptInjectionDetected, SecurityException

logger = structlog.get_logger(__name__)

# ── Prompt injection signatures ───────────────────────────────────────────────
# Common patterns used to override system instructions.
_INJECTION_PATTERNS: list[re.Pattern[str]] = [
    re.compile(p, re.IGNORECASE)
    for p in [
        r"ignore\s+(all\s+)?previous\s+instructions",
        r"disregard\s+",
        r"forget\s+your\s+(instructions|rules|system\s+prompt)",
        r"you\s+are\s+now\s+(a\s+)?(?:different|new|evil|unrestricted)",
        r"override\s+(your\s+)?(instructions|rules|constraints)",
        r"act\s+as\s+(?:if\s+you\s+(?:have\s+no|are\s+without)\s+restrictions)",
        r"do\s+not\s+follow\s+your\s+(rules|guidelines)",
        r"system\s*:\s*you\s+are",
        r"<\s*system\s*>",
        r"###\s*new\s+instructions",
        r"reveal\s+your\s+(system\s+)?prompt",
        r"print\s+your\s+(system\s+)?prompt",
        r"show\s+me\s+your\s+(system\s+)?prompt",
    ]
]

# Dangerous SQL keywords that should never appear in a SELECT-only context.
_DANGEROUS_KEYWORDS: frozenset[str] = frozenset(
    [
        "INSERT",
        "UPDATE",
        "DELETE",
        "DROP",
        "TRUNCATE",
        "ALTER",
        "CREATE",
        "REPLACE",
        "MERGE",
        "EXEC",
        "EXECUTE",
        "CALL",
        "GRANT",
        "REVOKE",
        "LOAD",
        "IMPORT",
        "COPY",
        "ATTACH",
        "DETACH",
    ]
)


class SQLValidator:
    """
    Security gate for all SQL generated by the LLM.

    Usage::

        validator = SQLValidator(settings)
        validator.validate(sql)           # raises SecurityException if unsafe
        validator.check_question(question) # raises PromptInjectionDetected
    """

    def __init__(self, settings: Settings) -> None:
        self._allowed_ops: frozenset[str] = frozenset(
            op.upper() for op in settings.allowed_sql_operations
        )
        self._injection_check: bool = settings.enable_injection_detection

    # ── Public API ────────────────────────────────────────────────────────────

    def validate(self, sql: str) -> None:
        """
        Full security validation of a SQL string.
        Raises SecurityException on any failure.
        Never returns a value — either passes or raises.
        """
        stripped = sql.strip()

        if not stripped:
            raise SecurityException("Empty SQL statement received.", sql=sql)

        self._check_dangerous_keywords(stripped)
        self._check_allowed_operations(stripped)
        self._check_with_sqlglot(stripped)

        logger.debug("validator.passed", sql_preview=stripped[:60])

    def check_question(self, question: str) -> None:
        """
        Scan a natural language question for prompt injection attempts.
        Raises PromptInjectionDetected if a pattern matches.
        Call this before SQL generation, not after.
        """
        if not self._injection_check:
            return
        for pattern in _INJECTION_PATTERNS:
            if pattern.search(question):
                logger.warning(
                    "validator.injection_detected",
                    pattern=pattern.pattern,
                    question=question[:80],
                )
                raise PromptInjectionDetected(
                    reason=f"Prompt injection pattern detected: '{pattern.pattern}'",
                    sql="",
                )

    # ── Private checks ────────────────────────────────────────────────────────

    def _check_dangerous_keywords(self, sql: str) -> None:
        """
        Keyword scan — catches obvious dangerous statements even before parsing.
        We normalise whitespace and strip comments first.
        """
        normalised = self._strip_comments(sql).upper()
        # Tokenise on non-alphanumeric boundaries to avoid false positives
        # e.g. 'SELECTED' should not trigger 'DELETE'
        tokens: set[str] = set(re.findall(r"\b[A-Z_]+\b", normalised))
        bad = tokens & _DANGEROUS_KEYWORDS
        if bad:
            raise SecurityException(
                f"Disallowed SQL keyword(s) detected: {sorted(bad)}",
                sql=sql,
            )

    def _check_allowed_operations(self, sql: str) -> None:
        """
        Ensure the first meaningful keyword is in the allowed operations list.
        Handles WITH ... AS (...) SELECT ... (CTEs) correctly.
        """
        normalised = self._strip_comments(sql).strip().upper()
        first_keyword = re.split(r"\s+", normalised)[0] if normalised else ""

        if first_keyword not in self._allowed_ops:
            raise SecurityException(
                f"SQL operation '{first_keyword}' is not permitted. "
                f"Allowed: {sorted(self._allowed_ops)}",
                sql=sql,
            )

    def _check_with_sqlglot(self, sql: str) -> None:
        """
        Parse the SQL with sqlglot to catch structural issues and disguised
        dangerous statements that keyword scanning might miss (e.g. stacked
        queries using semicolons).
        """
        try:
            statements = sqlglot.parse(sql)
        except Exception:
            # sqlglot could not parse it at all — treat as suspicious
            raise SecurityException(
                "SQL could not be parsed — it may be malformed or use an " "unsupported dialect.",
                sql=sql,
            ) from None

        if not statements:
            raise SecurityException("SQL parsed to zero statements.", sql=sql)

        # Guard against multiple statements (e.g. SELECT 1; DROP TABLE users)
        if len(statements) > 1:
            raise SecurityException(
                f"Multiple SQL statements detected ({len(statements)}). "
                "Only a single SELECT is permitted.",
                sql=sql,
            )

        stmt = statements[0]
        if stmt is None:
            raise SecurityException("SQL parsed to a null statement.", sql=sql)

        # Check the AST node type
        stmt_type = type(stmt).__name__.upper()
        allowed_ast_types = {"SELECT", "UNION", "INTERSECT", "EXCEPT", "WITH"}
        if stmt_type not in allowed_ast_types:
            raise SecurityException(
                f"Parsed statement type '{stmt_type}' is not allowed.",
                sql=sql,
            )

    @staticmethod
    def _strip_comments(sql: str) -> str:
        """Remove -- line comments and /* block comments */ from SQL."""
        # Remove block comments
        sql = re.sub(r"/\*.*?\*/", " ", sql, flags=re.DOTALL)
        # Remove line comments
        sql = re.sub(r"--[^\n]*", " ", sql)
        return sql
