"""SQL database adapters behind the one "database" Connector type.

Task Designer's "SQL — Query" / "SQL — Execute" nodes run one statement
against an external database. SQL is always written with `:name`
placeholders, bound from the node's declared inputs as real driver
parameters — values are never pasted into the SQL text — and each driver's
own placeholder style is produced by bind_named():

    PostgreSQL      asyncpg    $1, $2 …
    MySQL/MariaDB   PyMySQL    %(name)s
    SQL Server      python-tds %(name)s
    Oracle          oracledb   :name (native)

Queries always run read-only: statements are checked (one statement, no
write keywords) and, where the database supports it, inside a READ ONLY
transaction. Writes need a connector with allow_writes = "yes". The real
safety net is still the database login's own permissions — the connector's
help says so.

Every adapter opens one connection per call and closes it; there is no pool.
"""

import asyncio
import base64
import datetime as dt
import re
import uuid
from dataclasses import dataclass, field
from decimal import Decimal
from typing import Any

MAX_ROWS_LIMIT = 10_000
MAX_TIMEOUT = 300


class SqlError(Exception):
    pass


@dataclass
class SqlResult:
    columns: list[str] = field(default_factory=list)
    rows: list[dict] = field(default_factory=list)
    affected_rows: int | None = None
    truncated: bool = False


# ---- SQL text handling -------------------------------------------------------------------------

_IDENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")


def _scan(sql: str):
    """Yields (kind, text) chunks: "code", "string" (quoted literal or
    quoted identifier) or "comment" — enough of a lexer to find
    placeholders, statement separators and keywords outside literals."""
    i, n = 0, len(sql)
    start = 0
    while i < n:
        ch = sql[i]
        two = sql[i : i + 2]
        if ch in ("'", '"', "`"):
            if start < i:
                yield "code", sql[start:i]
            j = i + 1
            while j < n:
                if sql[j] == ch:
                    if j + 1 < n and sql[j + 1] == ch:  # doubled quote = escaped
                        j += 2
                        continue
                    break
                if sql[j] == "\\" and ch != '"':  # MySQL-style backslash escape
                    j += 2
                    continue
                j += 1
            yield "string", sql[i : j + 1]
            i = start = j + 1
            continue
        if two == "--":
            if start < i:
                yield "code", sql[start:i]
            j = sql.find("\n", i)
            j = n if j == -1 else j
            yield "comment", sql[i:j]
            i = start = j
            continue
        if two == "/*":
            if start < i:
                yield "code", sql[start:i]
            j = sql.find("*/", i + 2)
            j = n if j == -1 else j + 2
            yield "comment", sql[i:j]
            i = start = j
            continue
        if ch == "$" and (m := re.match(r"\$([A-Za-z_]*)\$", sql[i:])):
            # PostgreSQL dollar-quoted string: $$…$$ or $tag$…$tag$
            tag = m.group(0)
            end = sql.find(tag, i + len(tag))
            end = n if end == -1 else end + len(tag)
            if start < i:
                yield "code", sql[start:i]
            yield "string", sql[i:end]
            i = start = end
            continue
        i += 1
    if start < n:
        yield "code", sql[start:]


def placeholders(sql: str) -> list[str]:
    """The distinct :name placeholders, in first-use order."""
    names: list[str] = []
    for kind, text in _scan(sql):
        if kind != "code":
            continue
        for m in re.finditer(r"(?<![:\w]):([A-Za-z_][A-Za-z0-9_]*)", text):
            if m.group(1) not in names:
                names.append(m.group(1))
    return names


def bind_named(sql: str, params: dict, style: str) -> tuple[str, Any]:
    """Rewrites :name placeholders into the driver's style and returns
    (sql, args) — a list for "numeric" ($1), a dict otherwise. "::" casts and
    anything inside literals/comments are left alone."""
    missing = [n for n in placeholders(sql) if n not in params]
    if missing:
        raise SqlError(f"the SQL uses {', '.join(':' + m for m in missing)} but the node has no input with that name")
    order: list[str] = []
    out: list[str] = []

    def repl(m: re.Match) -> str:
        name = m.group(1)
        if style == "numeric":
            if name not in order:
                order.append(name)
            return f"${order.index(name) + 1}"
        if style == "pyformat":
            return f"%({name})s"
        return f":{name}"

    for kind, text in _scan(sql):
        if style == "pyformat":
            text = text.replace("%", "%%")  # the driver %-formats the whole statement
        if kind == "code":
            text = re.sub(r"(?<![:\w]):([A-Za-z_][A-Za-z0-9_]*)", repl, text)
        out.append(text)
    used = placeholders(sql)
    if style == "numeric":
        return "".join(out), [params[n] for n in order]
    return "".join(out), {n: params[n] for n in used}


_WRITE_WORDS = {
    "insert", "update", "delete", "merge", "upsert", "replace", "drop", "alter", "create", "truncate", "rename",
    "grant", "revoke", "exec", "execute", "call", "copy", "lock", "vacuum", "into", "set", "commit", "rollback",
}
_READ_STARTS = {"select", "with", "show", "explain", "describe", "desc", "values", "table"}


def check_statement(sql: str, *, read_only: bool) -> None:
    """One statement only; for read-only use, it must start like a query and
    mention no write keyword (outside literals/comments)."""
    code = " ".join(text for kind, text in _scan(sql) if kind == "code")
    body = code.strip().rstrip(";").strip()
    if not body:
        raise SqlError("no SQL to run")
    if ";" in body:
        raise SqlError("one statement per node — split it into several nodes")
    if not read_only:
        return
    words = [w.lower() for w in _IDENT.findall(body)]
    if not words or words[0] not in _READ_STARTS:
        raise SqlError("SQL — Query only runs SELECT statements; use SQL — Execute to change data")
    # SET and REPLACE() are harmless inside a query; UPDATE…SET is caught by UPDATE.
    found = sorted({w for w in words if w in _WRITE_WORDS} - {"set", "replace"})
    if found:
        raise SqlError(f"a read-only query can't use {', '.join(w.upper() for w in found)} — use SQL — Execute on a connector that allows writes")


def jsonable(value: Any) -> Any:
    """Driver values -> what a task context (JSON) can hold."""
    if value is None or isinstance(value, (bool, int, float, str)):
        return value
    if isinstance(value, Decimal):
        return int(value) if value == value.to_integral_value() and abs(value) < 2**63 else float(value)
    if isinstance(value, (dt.datetime, dt.date, dt.time)):
        return value.isoformat()
    if isinstance(value, dt.timedelta):
        return value.total_seconds()
    if isinstance(value, uuid.UUID):
        return str(value)
    if isinstance(value, (bytes, bytearray, memoryview)):
        return base64.b64encode(bytes(value)).decode()
    if isinstance(value, (list, tuple)):
        return [jsonable(v) for v in value]
    if isinstance(value, dict):
        return {str(k): jsonable(v) for k, v in value.items()}
    if hasattr(value, "read"):  # Oracle LOB
        return jsonable(value.read())
    return str(value)


def _rows(columns: list[str], raw: list, max_rows: int) -> tuple[list[dict], bool]:
    truncated = len(raw) > max_rows
    return [{c: jsonable(v) for c, v in zip(columns, r)} for r in raw[:max_rows]], truncated


def _settings(config: dict, default_port: int) -> tuple[str, int, str, str, str]:
    host = (config.get("host") or "").strip()
    database = (config.get("database") or "").strip()
    user = (config.get("username") or "").strip()
    if not host or not user:
        raise SqlError("the connector needs the host and username")
    try:
        port = int(config.get("port") or default_port)
    except ValueError as exc:
        raise SqlError("the port must be a number") from exc
    return host, port, database, user, config.get("password") or ""


def _encrypt(config: dict) -> bool:
    return (config.get("ssl") or "prefer") == "require"


# ---- PostgreSQL ---------------------------------------------------------------------------------


class Postgres:
    name = "PostgreSQL"

    async def _connect(self, config: dict, timeout: float):
        import asyncpg

        host, port, database, user, password = _settings(config, 5432)
        return await asyncpg.connect(
            host=host, port=port, database=database or None, user=user, password=password,
            ssl=config.get("ssl") if config.get("ssl") in ("require", "disable") else "prefer",
            timeout=min(timeout, 15), command_timeout=timeout,
        )

    async def run(self, config: dict, sql: str, params: dict, *, read_only: bool, max_rows: int, timeout: float) -> SqlResult:
        text, args = bind_named(sql, params, "numeric")
        conn = await self._connect(config, timeout)
        try:
            async with conn.transaction(readonly=read_only):
                stmt = await conn.prepare(text)
                columns = [a.name for a in stmt.get_attributes()]
                if columns:
                    cursor = await stmt.cursor(*args)
                    raw = await cursor.fetch(max_rows + 1)
                    rows, truncated = _rows(columns, [tuple(r) for r in raw], max_rows)
                else:
                    await stmt.fetch(*args)
                    rows, truncated = [], False
                status = stmt.get_statusmsg() or ""
            count = int(status.rsplit(" ", 1)[-1]) if status and status.rsplit(" ", 1)[-1].isdigit() else None
            return SqlResult(columns=columns, rows=rows, affected_rows=count, truncated=truncated)
        finally:
            await conn.close()

    async def describe(self, config: dict) -> str:
        conn = await self._connect(config, 15)
        try:
            version = await conn.fetchval("SHOW server_version")
            database = await conn.fetchval("SELECT current_database()")
            return f"PostgreSQL {version}, database {database}"
        finally:
            await conn.close()


# ---- DB-API drivers (MySQL, SQL Server, Oracle), run in a worker thread -------------------------


class _DbApi:
    name = ""
    style = "pyformat"
    default_port = 0
    read_only_sql: str | None = None
    version_sql = ""

    def connect(self, config: dict, timeout: float):  # pragma: no cover - per driver
        raise NotImplementedError

    def _run_sync(self, config: dict, sql: str, params: dict, read_only: bool, max_rows: int, timeout: float) -> SqlResult:
        text, args = bind_named(sql, params, self.style)
        conn = self.connect(config, timeout)
        try:
            cur = conn.cursor()
            if read_only and self.read_only_sql:
                cur.execute(self.read_only_sql)
            if args:
                cur.execute(text, args)
            else:  # without args the driver doesn't %-format, so undo the escaping
                cur.execute(text.replace("%%", "%") if self.style == "pyformat" else text)
            columns = [d[0] for d in cur.description] if cur.description else []
            rows, truncated = _rows(columns, list(cur.fetchmany(max_rows + 1)), max_rows) if columns else ([], False)
            count = cur.rowcount if isinstance(cur.rowcount, int) and cur.rowcount >= 0 else None
            if read_only:
                conn.rollback()
            else:
                conn.commit()
            return SqlResult(columns=columns, rows=rows, affected_rows=count, truncated=truncated)
        finally:
            conn.close()

    async def run(self, config: dict, sql: str, params: dict, *, read_only: bool, max_rows: int, timeout: float) -> SqlResult:
        return await asyncio.wait_for(asyncio.to_thread(self._run_sync, config, sql, params, read_only, max_rows, timeout), timeout=timeout + 15)

    def _describe_sync(self, config: dict) -> str:
        conn = self.connect(config, 15)
        try:
            cur = conn.cursor()
            cur.execute(self.version_sql)
            version = str(cur.fetchone()[0]).splitlines()[0]
            return f"{self.name} {version}"
        finally:
            conn.close()

    async def describe(self, config: dict) -> str:
        return await asyncio.wait_for(asyncio.to_thread(self._describe_sync, config), timeout=30)


class MySql(_DbApi):
    name = "MySQL"
    default_port = 3306
    read_only_sql = "START TRANSACTION READ ONLY"
    version_sql = "SELECT VERSION()"

    def connect(self, config: dict, timeout: float):
        import pymysql

        host, port, database, user, password = _settings(config, self.default_port)
        return pymysql.connect(
            host=host, port=port, user=user, password=password, database=database or None, charset="utf8mb4",
            connect_timeout=15, read_timeout=int(timeout), write_timeout=int(timeout), autocommit=False,
            ssl={"check_hostname": False} if _encrypt(config) else None,
        )


class SqlServer(_DbApi):
    # No READ ONLY transaction here — the statement check plus the rollback
    # at the end keep a query from changing anything.
    name = "SQL Server"
    default_port = 1433
    version_sql = "SELECT @@VERSION"

    def connect(self, config: dict, timeout: float):
        import certifi
        import pytds

        host, port, database, user, password = _settings(config, self.default_port)
        kwargs: dict = {"cafile": certifi.where(), "validate_host": True} if _encrypt(config) else {}
        return pytds.connect(dsn=host, port=port, database=database or None, user=user, password=password, login_timeout=15, timeout=timeout, autocommit=False, **kwargs)


class Oracle(_DbApi):
    name = "Oracle"
    style = "named"
    default_port = 1521
    read_only_sql = "SET TRANSACTION READ ONLY"
    version_sql = "SELECT banner FROM v$version WHERE ROWNUM = 1"

    def connect(self, config: dict, timeout: float):
        import oracledb

        host, port, service, user, password = _settings(config, self.default_port)
        if not service:
            raise SqlError("Oracle needs the service name (e.g. ORCLPDB1 or FREEPDB1)")
        dsn = f"{'tcps' if _encrypt(config) else 'tcp'}://{host}:{port}/{service}"
        conn = oracledb.connect(user=user, password=password, dsn=dsn, tcp_connect_timeout=15)
        conn.call_timeout = int(timeout * 1000)
        return conn

    def _describe_sync(self, config: dict) -> str:
        conn = self.connect(config, 15)
        try:
            return f"Oracle {conn.version}"
        finally:
            conn.close()


ADAPTERS: dict[str, Any] = {"postgresql": Postgres(), "mysql": MySql(), "sqlserver": SqlServer(), "oracle": Oracle()}


def adapter_for(config: dict):
    provider = config.get("provider")
    if provider not in ADAPTERS:
        raise SqlError("choose the database type on the connector")
    return ADAPTERS[provider]


def writes_allowed(config: dict) -> bool:
    return (config.get("allow_writes") or "no") == "yes"


async def run_sql(config: dict, sql: str, params: dict, *, write: bool, max_rows: int = 1000, timeout: float = 30) -> SqlResult:
    if write and not writes_allowed(config):
        raise SqlError("this connector is read-only — edit it and allow INSERT / UPDATE / DELETE to use SQL — Execute")
    check_statement(sql, read_only=not write)
    max_rows = max(1, min(int(max_rows or 1000), MAX_ROWS_LIMIT))
    timeout = float(max(1, min(int(timeout or 30), MAX_TIMEOUT)))
    return await adapter_for(config).run(config, sql, params, read_only=not write, max_rows=max_rows, timeout=timeout)


async def test_credentials(config: dict) -> tuple[bool, str]:
    described = await adapter_for(config).describe(config)
    mode = "writes allowed" if writes_allowed(config) else "read-only"
    return True, f"Connected to {described} ({mode})"


def _f(key: str, label: str, *, secret: bool = False, required: bool = True, placeholder: str = "", help: str = "", options: list | None = None, advanced: bool = False, multiline: bool = False) -> dict:
    return {"key": key, "label": label, "secret": secret, "required": required, "placeholder": placeholder, "help": help, "options": options, "advanced": advanced, "multiline": multiline}


def _common(port: int, database_label: str, database_placeholder: str, database_required: bool = False) -> list[dict]:
    return [
        _f("host", "Host", placeholder="db.example.com"),
        _f("port", "Port", required=False, placeholder=str(port)),
        _f("database", database_label, required=database_required, placeholder=database_placeholder),
        _f("username", "Username", help="Use a login with only the permissions your tasks need — read-only unless they must write."),
        _f("password", "Password", secret=True, required=False),
        _f("ssl", "Encryption", options=[["prefer", "Use TLS when the server offers it"], ["require", "Require TLS"], ["disable", "Off"]]),
        _f("allow_writes", "Access", options=[["no", "Read-only — queries only (recommended)"], ["yes", "Allow INSERT / UPDATE / DELETE"]], help="Read-only connectors can only be used by SQL — Query nodes."),
    ]


PROVIDERS: list[dict] = [
    {"id": "postgresql", "name": "PostgreSQL", "tagline": "Also Amazon RDS / Aurora, Azure and Cloud SQL for PostgreSQL", "environments": [], "docs_url": None, "notes": "", "capabilities": [], "fields": _common(5432, "Database", "postgres")},
    {"id": "mysql", "name": "MySQL / MariaDB", "tagline": "Also Amazon RDS / Aurora MySQL and Azure Database for MySQL", "environments": [], "docs_url": None, "notes": "", "capabilities": [], "fields": _common(3306, "Database", "app")},
    {"id": "sqlserver", "name": "SQL Server", "tagline": "Microsoft SQL Server and Azure SQL (SQL logins)", "environments": [], "docs_url": None, "notes": "Azure SQL requires TLS — choose \"Require TLS\".", "capabilities": [], "fields": _common(1433, "Database", "master")},
    {"id": "oracle", "name": "Oracle", "tagline": "Oracle Database 12.1+ and Autonomous DB (no Instant Client needed)", "environments": [], "docs_url": None, "notes": "", "capabilities": [], "fields": _common(1521, "Service name", "ORCLPDB1", database_required=True)},
]