# retoor <retoor@molodetz.nl>

from __future__ import annotations

import re
import sqlite3
from dataclasses import dataclass, field

import sqlglot
from sqlglot import exp

from devplacepy.config import DATABASE_URL

UNSAFE = re.compile(r"\b(attach|detach|pragma|vacuum|reindex)\b", re.IGNORECASE)

_STATEMENT_TYPES = {
    exp.Select: "select",
    exp.Union: "select",
    exp.Insert: "insert",
    exp.Update: "update",
    exp.Delete: "delete",
    exp.Create: "ddl",
    exp.Drop: "ddl",
    exp.Alter: "ddl",
}


@dataclass
class Verdict:
    sql: str
    statement_type: str = "other"
    is_select: bool = False
    valid: bool = False
    tables: list[str] = field(default_factory=list)
    has_where: bool = False
    has_join: bool = False
    has_limit: bool = False
    suspicious: list[str] = field(default_factory=list)
    error: str | None = None

    def as_dict(self) -> dict:
        return {
            "sql": self.sql,
            "statement_type": self.statement_type,
            "is_select": self.is_select,
            "valid": self.valid,
            "tables": self.tables,
            "has_where": self.has_where,
            "has_join": self.has_join,
            "has_limit": self.has_limit,
            "suspicious": self.suspicious,
            "error": self.error,
        }


def _database_path() -> str:
    return DATABASE_URL.split("sqlite:///", 1)[-1]


def _statement_type(node: exp.Expression) -> str:
    for kind, label in _STATEMENT_TYPES.items():
        if isinstance(node, kind):
            return label
    return "other"


def classify(sql: str, dialect: str = "sqlite") -> Verdict:
    text = (sql or "").strip().rstrip(";").strip()
    verdict = Verdict(sql=text)
    if not text:
        verdict.error = "Empty query."
        return verdict
    try:
        statements = [node for node in sqlglot.parse(text, read=dialect) if node]
    except Exception as exc:
        verdict.error = f"Could not parse SQL: {str(exc).splitlines()[0][:200]}"
        return verdict
    if not statements:
        verdict.error = "No statement found."
        return verdict
    node = statements[0]
    verdict.statement_type = _statement_type(node)
    verdict.is_select = verdict.statement_type == "select"
    verdict.tables = sorted({t.name for t in node.find_all(exp.Table) if t.name})
    verdict.has_where = node.find(exp.Where) is not None
    verdict.has_join = node.find(exp.Join) is not None
    verdict.has_limit = node.find(exp.Limit) is not None
    if len(statements) > 1:
        verdict.suspicious.append("Multiple statements in one query are not allowed.")
    if UNSAFE.search(text):
        verdict.suspicious.append("ATTACH/PRAGMA/VACUUM style statements are not allowed.")
        verdict.is_select = False
    if verdict.is_select and not (
        verdict.has_where or verdict.has_join or verdict.has_limit
    ):
        verdict.suspicious.append(
            "SELECT has no WHERE, JOIN, or LIMIT and may return an entire table."
        )
    return verdict


def dry_run(sql: str) -> tuple[bool, str | None]:
    text = (sql or "").strip().rstrip(";").strip()
    connection = sqlite3.connect(f"file:{_database_path()}?mode=ro", uri=True, timeout=5)
    try:
        connection.execute("PRAGMA query_only=ON")
        connection.execute(f"EXPLAIN {text}")
        return True, None
    except sqlite3.Error as exc:
        return False, str(exc)[:300]
    finally:
        connection.close()


def validate_select(sql: str, dialect: str = "sqlite") -> Verdict:
    verdict = classify(sql, dialect=dialect)
    if verdict.error:
        return verdict
    if len(verdict.suspicious) and any(
        "Multiple statements" in note or "not allowed" in note
        for note in verdict.suspicious
    ):
        verdict.error = verdict.suspicious[0]
        return verdict
    if not verdict.is_select:
        verdict.error = (
            f"Only SELECT queries run through query(); this is a {verdict.statement_type} "
            "statement. The database API is read-only and cannot change data."
        )
        return verdict
    ok, error = dry_run(verdict.sql)
    verdict.valid = ok
    if not ok:
        verdict.error = error
    return verdict


def run_select(sql: str, limit: int) -> tuple[list[dict], bool]:
    text = sql.strip().rstrip(";").strip()
    connection = sqlite3.connect(f"file:{_database_path()}?mode=ro", uri=True, timeout=30)
    try:
        connection.row_factory = sqlite3.Row
        connection.execute("PRAGMA query_only=ON")
        cursor = connection.execute(text)
        rows = [dict(row) for row in cursor.fetchmany(limit + 1)]
        truncated = len(rows) > limit
        return rows[:limit], truncated
    finally:
        connection.close()


def transpile(sql: str, read: str = "sqlite", write: str = "sqlite") -> str:
    try:
        return sqlglot.transpile(sql, read=read, write=write)[0]
    except Exception:
        return sql
