from __future__ import annotations
import re
from typing import Iterable
from .errors import QueryRejected

def normalize_sql(sql: str) -> str:
    return sql.strip().rstrip(";")

def reject_dangerous(sql: str, deny_patterns: Iterable[str]) -> None:
    for pat in deny_patterns:
        if re.search(pat, sql, flags=re.IGNORECASE):
            raise QueryRejected(f"Query rejected by policy: matched pattern: {pat}")

def ensure_limit(sql: str, max_rows: int) -> str:
    # If LIMIT exists, keep it; otherwise append LIMIT.
    if re.search(r"\bLIMIT\b", sql, flags=re.IGNORECASE):
        return sql
    return f"{sql}\nLIMIT {int(max_rows)}"

def enforce_guardrails(sql: str, deny_patterns: Iterable[str], max_rows: int) -> str:
    s = normalize_sql(sql)
    reject_dangerous(s, deny_patterns)
    s = ensure_limit(s, max_rows)
    return s

