#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Dependency-free Spark SQL allowlist guard.

The guard answers one question only: is this SQL statement type safe enough
to show or pass to spark.sql? It does not prove business correctness.
"""

from __future__ import annotations

import re
from dataclasses import dataclass

ALLOWED_FIRST = {"select", "with", "insert"}
FORBIDDEN = {
    "delete", "truncate", "drop", "alter", "create", "replace", "merge", "update",
    "msck", "refresh", "analyze", "cache", "uncache", "vacuum", "optimize",
    "call", "grant", "revoke",
}


@dataclass(frozen=True)
class GuardResult:
    status: str
    first_keyword: str
    reason: str = ""

    def as_dict(self) -> dict[str, str]:
        return {
            "status": self.status,
            "first_keyword": self.first_keyword,
            "reason": self.reason,
        }


def _mask(sql: str) -> str:
    """Mask strings, comments, and backtick identifiers before token checks."""
    out: list[str] = []
    i = 0
    state = "normal"
    quote = ""
    while i < len(sql):
        ch = sql[i]
        nxt = sql[i + 1] if i + 1 < len(sql) else ""
        if state == "normal":
            if ch in {"'", '"', "`"}:
                state = "quote"
                quote = ch
                out.append(" ")
                i += 1
            elif ch == "-" and nxt == "-":
                state = "line_comment"
                out.extend("  ")
                i += 2
            elif ch == "/" and nxt == "*":
                state = "block_comment"
                out.extend("  ")
                i += 2
            else:
                out.append(ch)
                i += 1
        elif state == "quote":
            out.append(" ")
            if ch == quote:
                if nxt == quote and quote in {"'", '"'}:
                    out.append(" ")
                    i += 2
                else:
                    state = "normal"
                    i += 1
            elif ch == "\\" and nxt:
                out.append(" ")
                i += 2
            else:
                i += 1
        elif state == "line_comment":
            out.append("\n" if ch == "\n" else " ")
            if ch == "\n":
                state = "normal"
            i += 1
        else:
            out.append("\n" if ch == "\n" else " ")
            if ch == "*" and nxt == "/":
                out.append(" ")
                state = "normal"
                i += 2
            else:
                i += 1
    return "".join(out)


def _body(masked: str) -> str:
    text = masked.strip()
    return text[:-1].strip() if text.endswith(";") else text


def first_keyword(sql: str) -> str:
    match = re.search(r"\b([A-Za-z_][\w]*)\b", _body(_mask(sql)))
    return match.group(1).lower() if match else "unknown"


def validate_sql(sql: str, *, allow_overwrite: bool = False) -> GuardResult:
    if not sql or not sql.strip():
        return GuardResult("FAIL", "unknown", "SQL is empty")

    masked_body = _body(_mask(sql))
    first = first_keyword(sql)

    if ";" in masked_body:
        return GuardResult("FAIL", first, "Multiple SQL statements are not allowed")
    if first not in ALLOWED_FIRST:
        return GuardResult("FAIL", first, f"Only SELECT, WITH, or INSERT is allowed; got {first!r}")

    lowered = masked_body.lower()
    found = sorted(word for word in FORBIDDEN if re.search(rf"\b{word}\b", lowered))
    if found:
        return GuardResult("FAIL", first, "Forbidden keyword found: " + ", ".join(found))
    if first == "with" and not re.search(r"\bselect\b", lowered):
        return GuardResult("FAIL", first, "WITH must contain SELECT")
    if first == "insert":
        if not re.search(r"\bselect\b", lowered):
            return GuardResult("FAIL", first, "INSERT must be based on SELECT")
        if re.search(r"\binsert\s+overwrite\b", lowered) and not allow_overwrite:
            return GuardResult("FAIL", first, "INSERT OVERWRITE requires --allow-overwrite")

    return GuardResult("PASS", first, "")


def assert_allowed_sql(sql: str, *, allow_overwrite: bool = False) -> str:
    result = validate_sql(sql, allow_overwrite=allow_overwrite)
    if result.status != "PASS":
        raise ValueError(result.reason)
    return sql.strip().rstrip(";")


def safe_spark_sql(spark, sql: str, *, allow_overwrite: bool = False):
    return spark.sql(assert_allowed_sql(sql, allow_overwrite=allow_overwrite))