"""SQL AST Analyzer — parse SQL using sqlparse for safe query validation.

Used by query endpoint to determine if a SQL statement is safe for read-only execution.
Separate from audit_engine.py (which is for ticket review scoring).
"""
import sqlparse
from sqlparse.sql import IdentifierList, Identifier, Where, Comparison
from sqlparse.tokens import Keyword, DML, DDL, Punctuation
from dataclasses import dataclass
from typing import List, Optional
from enum import Enum


class SQLType(str, Enum):
    SELECT = "SELECT"
    INSERT = "INSERT"
    UPDATE = "UPDATE"
    DELETE = "DELETE"
    CREATE = "CREATE"
    ALTER = "ALTER"
    DROP = "DROP"
    TRUNCATE = "TRUNCATE"
    UNKNOWN = "UNKNOWN"


class RiskLevel(str, Enum):
    SAFE = "safe"           # SELECT only
    LOW = "low"             # SELECT with functions
    MEDIUM = "medium"       # DML with WHERE
    HIGH = "high"           # DML without WHERE
    DANGEROUS = "dangerous"  # DDL operations


@dataclass
class SQLAnalysis:
    sql_type: SQLType
    risk_level: RiskLevel
    is_readonly: bool
    has_where: bool
    tables: List[str]
    warnings: List[str]
    errors: List[str]

    def to_dict(self) -> dict:
        return {
            "sql_type": self.sql_type.value,
            "risk_level": self.risk_level.value,
            "is_readonly": self.is_readonly,
            "has_where": self.has_where,
            "tables": self.tables,
            "warnings": self.warnings,
            "errors": self.errors,
        }


class SQLAnalyzer:
    """Parse and analyze SQL statements using sqlparse AST."""

    # Statements that are always read-only
    READONLY_TYPES = {SQLType.SELECT}

    # Statements that modify data
    DML_TYPES = {SQLType.INSERT, SQLType.UPDATE, SQLType.DELETE}

    # Statements that change schema
    DDL_TYPES = {SQLType.CREATE, SQLType.ALTER, SQLType.DROP, SQLType.TRUNCATE}

    def analyze(self, sql: str) -> SQLAnalysis:
        """Analyze a SQL statement and return detailed analysis."""
        sql = sql.strip()

        # Detect SQL type from the first meaningful token
        sql_type = self._detect_sql_type(sql)

        # Parse with sqlparse
        parsed = sqlparse.parse(sql)
        if not parsed:
            return SQLAnalysis(
                sql_type=SQLType.UNKNOWN, risk_level=RiskLevel.SAFE,
                is_readonly=True, has_where=False, tables=[],
                warnings=[], errors=["无法解析 SQL 语句"],
            )

        statement = parsed[0]
        tokens = list(statement.flatten())

        # Extract tables
        tables = self._extract_tables(statement)

        # Check for WHERE clause
        has_where = self._has_where(statement)

        # Check for multiple statements (SQL injection vector)
        multi_errors = self._check_multi_statement(sql)

        # Determine risk level and generate warnings/errors
        warnings = []
        errors = list(multi_errors)

        if sql_type in self.DDL_TYPES:
            risk_level = RiskLevel.DANGEROUS
            is_readonly = False
            errors.append(f"DDL 操作 ({sql_type.value}) 在查询模式下被禁止")

        elif sql_type in self.DML_TYPES:
            is_readonly = False
            if has_where:
                risk_level = RiskLevel.MEDIUM
                warnings.append(f"{sql_type.value} 操作包含 WHERE 条件，但查询模式为只读")
            else:
                risk_level = RiskLevel.HIGH
                errors.append(f"{sql_type.value} 缺少 WHERE 条件，影响全表")
            errors.append(f"查询模式禁止执行 {sql_type.value} 操作")

        elif sql_type == SQLType.SELECT:
            is_readonly = True
            risk_level = RiskLevel.SAFE
            # Check for INTO OUTFILE (writes to filesystem)
            if self._has_into_outfile(statement):
                risk_level = RiskLevel.DANGEROUS
                is_readonly = False
                errors.append("SELECT INTO OUTFILE 被禁止（文件写入风险）")

            # Check for subqueries that modify data
            sub_type = self._check_subquery_dml(sql)
            if sub_type:
                is_readonly = False
                risk_level = RiskLevel.DANGEROUS
                errors.append(f"子查询包含 {sub_type} 操作，被禁止")
        else:
            is_readonly = True
            risk_level = RiskLevel.LOW

        return SQLAnalysis(
            sql_type=sql_type,
            risk_level=risk_level,
            is_readonly=is_readonly,
            has_where=has_where,
            tables=tables,
            warnings=warnings,
            errors=errors,
        )

    def validate_readonly(self, sql: str) -> tuple[bool, str]:
        """Validate SQL is read-only safe. Returns (is_safe, error_message)."""
        analysis = self.analyze(sql)
        if analysis.errors:
            return False, "; ".join(analysis.errors)
        if not analysis.is_readonly:
            return False, f"查询模式仅允许 SELECT 语句，当前为 {analysis.sql_type.value}"
        return True, ""

    def _detect_sql_type(self, sql: str) -> SQLType:
        """Detect the SQL statement type from the first keyword."""
        upper = sql.strip().upper()
        type_map = {
            "SELECT": SQLType.SELECT,
            "INSERT": SQLType.INSERT,
            "UPDATE": SQLType.UPDATE,
            "DELETE": SQLType.DELETE,
            "CREATE": SQLType.CREATE,
            "ALTER": SQLType.ALTER,
            "DROP": SQLType.DROP,
            "TRUNCATE": SQLType.TRUNCATE,
            "WITH": SQLType.SELECT,  # CTE is still SELECT
            "EXPLAIN": SQLType.SELECT,
        }
        for keyword, sql_type in type_map.items():
            if upper.startswith(keyword):
                return sql_type
        return SQLType.UNKNOWN

    def _extract_tables(self, statement) -> List[str]:
        """Extract table names from the statement."""
        tables = []
        seen = set()

        # Get tokens after FROM or JOIN or INTO or UPDATE
        tokens = list(statement.tokens)
        capture = False
        for token in tokens:
            if token.ttype in (Keyword,) and token.normalized in ("FROM", "INTO", "UPDATE", "JOIN", "TABLE"):
                capture = True
                continue
            if capture:
                if isinstance(token, IdentifierList):
                    for ident in token.get_identifiers():
                        name = self._get_identifier_name(ident)
                        if name and name not in seen:
                            seen.add(name)
                            tables.append(name)
                elif isinstance(token, Identifier):
                    name = self._get_identifier_name(token)
                    if name and name not in seen:
                        seen.add(name)
                        tables.append(name)
                capture = False

        return tables

    def _get_identifier_name(self, ident) -> Optional[str]:
        """Get the table name from an identifier."""
        if isinstance(ident, Identifier):
            # Use the real name (not alias)
            return ident.get_real_name()
        return None

    def _has_where(self, statement) -> bool:
        """Check if the statement has a WHERE clause."""
        for token in statement.tokens:
            if isinstance(token, Where):
                return True
        # Also check flattened tokens
        for token in statement.flatten():
            if token.ttype is Keyword and token.normalized == "WHERE":
                return True
        return False

    def _has_into_outfile(self, statement) -> bool:
        """Check for INTO OUTFILE (filesystem write)."""
        upper = str(statement).upper()
        return "INTO OUTFILE" in upper or "INTO DUMPFILE" in upper

    def _check_multi_statement(self, sql: str) -> List[str]:
        """Check for multiple statements (semicolons inside the SQL)."""
        errors = []
        stripped = sql.strip().rstrip(";")
        # Count semicolons not inside strings
        in_string = False
        string_char = None
        semi_count = 0
        for ch in stripped:
            if ch in ("'", '"') and not in_string:
                in_string = True
                string_char = ch
            elif ch == string_char and in_string:
                in_string = False
                string_char = None
            elif ch == ";" and not in_string:
                semi_count += 1

        if semi_count > 0:
            errors.append("查询不允许包含多条 SQL 语句（分号分隔）")
        return errors

    def _check_subquery_dml(self, sql: str) -> Optional[str]:
        """Check if a SELECT contains DML subqueries."""
        upper = sql.upper()
        # Simple heuristic: look for INSERT/UPDATE/DELETE after a parenthesis
        for kw in ("INSERT", "UPDATE", "DELETE", "CREATE", "ALTER", "DROP"):
            # Check if keyword appears inside parentheses
            pattern = f"({kw} "
            if pattern in upper:
                return kw
        return None