"""Shared foundation for SQL database connectors (Postgres, MySQL, ...).

Holds the SQL safety guard and a ``BaseSqlConnector`` that implements the common
``query`` (read-only, guarded) and ``execute`` (single-statement, sensitive)
tools. Concrete connectors only supply the driver-specific bits: how to connect
and run a statement.
"""

from __future__ import annotations

import re

from legate.agent.tooling import ToolContext, ToolResult, ToolSpec
from legate.connectors.base import BaseConnector, ConnectionResult

_DEFAULT_MAX_ROWS = 100
_HARD_MAX_ROWS = 1000

# Data-modifying / structural keywords that must never appear in a read query.
_WRITE_KEYWORDS = re.compile(
    r"\b("
    r"insert|update|delete|drop|alter|create|truncate|grant|revoke|"
    r"copy|merge|call|do|vacuum|reindex|refresh|comment|lock|set"
    r")\b",
    re.IGNORECASE,
)
_COMMENT_BLOCK = re.compile(r"/\*.*?\*/", re.DOTALL)
_COMMENT_LINE = re.compile(r"--[^\n]*")


class SqlSafetyError(ValueError):
    """Raised when a statement fails the safety checks."""


def _strip_comments(sql: str) -> str:
    return _COMMENT_LINE.sub(" ", _COMMENT_BLOCK.sub(" ", sql)).strip()


def _single_statement(sql: str) -> str:
    """Return the statement, rejecting anything that chains multiple statements."""
    cleaned = _strip_comments(sql).rstrip(";").strip()
    if ";" in cleaned:
        raise SqlSafetyError("only a single SQL statement is allowed")
    if not cleaned:
        raise SqlSafetyError("empty SQL statement")
    return cleaned


def _ensure_read_only(sql: str) -> str:
    stmt = _single_statement(sql)
    lowered = stmt.lower()
    if not (lowered.startswith("select") or lowered.startswith("with")):
        raise SqlSafetyError("query must be a SELECT (or WITH ... SELECT) statement")
    if _WRITE_KEYWORDS.search(stmt):
        raise SqlSafetyError("query must not contain data-modifying keywords")
    return stmt


def _sql_config_schema(*, extra: dict | None = None) -> dict:
    props: dict = {
        "host": {"type": "string"},
        "port": {"type": "integer"},
        "database": {"type": "string"},
        "user": {"type": "string"},
        "max_rows": {"type": "integer", "description": f"Row cap (<= {_HARD_MAX_ROWS})."},
    }
    if extra:
        props.update(extra)
    return {"type": "object", "properties": props, "required": ["host", "database", "user"]}


class BaseSqlConnector(BaseConnector):
    """Common query/execute behaviour for SQL databases.

    Subclasses set ``type``/``title``/``_default_port`` and implement the two
    driver seams (``_run_query`` / ``_run_exec``); everything else — the tool
    specs, the safety guard, row caps, and the connection test — is shared.
    """

    _default_port = 0
    config_schema = _sql_config_schema()
    secret_schema = {
        "type": "object",
        "properties": {"password": {"type": "string"}},
        "required": ["password"],
    }

    def list_tools(self) -> list[ToolSpec]:
        return [
            ToolSpec(
                name=f"{self.type}.query",
                description="Run a read-only SELECT and return rows (capped).",
                parameters={
                    "type": "object",
                    "properties": {
                        "sql": {"type": "string", "description": "A single SELECT statement."},
                        "params": {"type": "array", "description": "Positional parameters."},
                        "max_rows": {"type": "integer"},
                    },
                    "required": ["sql"],
                },
                is_sensitive=False,
                scopes=[f"{self.type}:read"],
            ),
            ToolSpec(
                name=f"{self.type}.execute",
                description="Run a single data-modifying statement (INSERT/UPDATE/DELETE).",
                parameters={
                    "type": "object",
                    "properties": {
                        "sql": {"type": "string", "description": "A single statement."},
                        "params": {"type": "array", "description": "Positional parameters."},
                    },
                    "required": ["sql"],
                },
                is_sensitive=True,
                scopes=[f"{self.type}:write"],
            ),
        ]

    def _max_rows(self, config: dict, args: dict) -> int:
        requested = int(args.get("max_rows") or config.get("max_rows") or _DEFAULT_MAX_ROWS)
        return max(1, min(requested, _HARD_MAX_ROWS))

    # -- driver seams (stubbed in tests, overridden by subclasses) --------

    def _run_query(
        self, config: dict, secrets: dict, sql: str, params: list | None, max_rows: int
    ) -> dict:  # pragma: no cover - overridden
        raise NotImplementedError

    def _run_exec(
        self, config: dict, secrets: dict, sql: str, params: list | None
    ) -> int:  # pragma: no cover - overridden
        raise NotImplementedError

    # -- connector API ----------------------------------------------------

    async def test_connection(self, config: dict, secrets: dict) -> ConnectionResult:
        try:
            self._run_query(config, secrets, "SELECT 1", None, 1)
        except Exception as exc:
            return ConnectionResult(ok=False, message=f"connection failed: {exc}")
        return ConnectionResult(ok=True, message="connected")

    async def execute(
        self, tool: str, args: dict, config: dict, secrets: dict, ctx: ToolContext
    ) -> ToolResult:
        sql = str(args.get("sql", ""))
        params = args.get("params")

        if tool == f"{self.type}.query":
            try:
                safe_sql = _ensure_read_only(sql)
            except SqlSafetyError as exc:
                return ToolResult(ok=False, error=str(exc))
            try:
                result = self._run_query(
                    config, secrets, safe_sql, params, self._max_rows(config, args)
                )
            except Exception as exc:
                return ToolResult(ok=False, error=f"query failed: {exc}")
            return ToolResult(
                ok=True,
                data={
                    "columns": result["columns"],
                    "rows": result["rows"],
                    "row_count": len(result["rows"]),
                },
            )

        if tool == f"{self.type}.execute":
            try:
                safe_sql = _single_statement(sql)
            except SqlSafetyError as exc:
                return ToolResult(ok=False, error=str(exc))
            try:
                affected = self._run_exec(config, secrets, safe_sql, params)
            except Exception as exc:
                return ToolResult(ok=False, error=f"execute failed: {exc}")
            return ToolResult(ok=True, data={"rows_affected": affected})

        return ToolResult(ok=False, error=f"unknown {self.type} tool: {tool}")
