#!/usr/bin/env python3
"""
localDbProbe.py - one-shot, read-only database reachability probe run BY THE
RUNNER on behalf of the Melaya server.

Why it exists: some databases sit behind an IP allow-list or inside a VPC that
the Melaya cloud can never reach, but the user's own runner (on their LAN /
bastion) can. The server hands us the resolved credentials over the already-
authenticated Socket.IO channel; we open a single real connection from THIS
machine's network vantage point, run a trivial read-only probe, and print the
outcome. We never persist the credentials and never write to the database.

Contract:
  stdin  : one JSON object {"service": "...", "creds": { ... }}
  stdout : one JSON object {"ok": bool, "message"|"error": "..."}
Exit code is always 0 (the result is the JSON on stdout); any crash is caught
and reported as {"ok": false, "error": ...}.
"""
import sys
import json
import base64
import hashlib
import urllib.parse


def _out(ok, text):
    key = "message" if ok else "error"
    sys.stdout.write(json.dumps({"ok": bool(ok), key: str(text)[:300]}))
    sys.stdout.flush()
    sys.exit(0)


def _pick(creds, *keys):
    for k in keys:
        v = (creds.get(k) or "").strip()
        if v:
            return v
    return ""


_PLACEHOLDER_HOSTS = {"host", "hostname", "your-host", "your_host", "<host>", "myhost", "dbhost", "example.com"}


def _friendly_conn_error(e, host, port, engine):
    """Turn a raw driver/socket exception into an actionable message. The most
    common failure is a copy-pasted EXAMPLE DSN (host == 'host') or a Docker/VPC
    hostname that doesn't resolve from the runner's machine."""
    import socket
    msg = str(e) or e.__class__.__name__
    low = msg.lower()
    host = host or ""
    is_dns = isinstance(e, socket.gaierror) or "nodename nor servname" in low \
        or "name or service not known" in low or "getaddrinfo" in low \
        or "name does not resolve" in low or "no address associated" in low
    if is_dns:
        if host.lower() in _PLACEHOLDER_HOSTS:
            return (f"the host is literally \"{host}\" - that's the EXAMPLE DSN, not your database. "
                    f"Paste your real connection string with the actual host, e.g. an AWS RDS / Cloud SQL "
                    f"endpoint: {engine.lower()}://user:pass@your-db.abc123.us-east-1.rds.amazonaws.com:{port}/db.")
        if host and "." not in host and host not in ("localhost",):
            return (f"can't resolve host \"{host}\" from the runner's machine. A bare name like this is "
                    f"usually a Docker/compose service or a VPC-internal alias that only resolves elsewhere - "
                    f"use the host that's reachable from this machine (e.g. your RDS/Cloud SQL endpoint, or a VPN/bastion host).")
        return (f"can't resolve host \"{host}\" from the runner's machine - check the host in your DSN, and that "
                f"this machine can reach it (DNS + network route; for a private RDS/VPC DB the runner must sit in the VPC or on an allow-listed IP).")
    if isinstance(e, ConnectionRefusedError) or "refused" in low:
        return f"host \"{host}\" reachable but port {port} refused the connection - is {engine} listening there / bound to that interface?"
    if isinstance(e, (TimeoutError,)) or "timed out" in low or "timeout" in low:
        return f"timed out connecting to {host}:{port} - firewall, IP allow-list, or wrong host/port?"
    if "password authentication failed" in low or "authentication failed" in low or "access denied" in low:
        return "authentication failed - check the username/password in the DSN."
    if 'database "' in low and "does not exist" in low or "unknown database" in low:
        return f"connected, but the database name in the DSN doesn't exist: {msg}"
    return msg


# ── PostgreSQL (asyncpg) ─────────────────────────────────────────────────────
def probe_postgres(creds):
    dsn = _pick(creds, "dsn", "POSTGRES_DSN")
    if not dsn:
        _out(False, "Connection DSN required")
    try:
        import asyncio
        import asyncpg  # type: ignore
    except ImportError:
        _out(False, "asyncpg not installed on this runner (pip install asyncpg)")
    p = urllib.parse.urlparse(dsn)
    host, port = (p.hostname or ""), (p.port or 5432)
    if not host:
        _out(False, "PostgreSQL connection failed: no host in the DSN - expected postgresql://user:pass@HOST:5432/db")

    async def _run():
        conn = await asyncpg.connect(dsn, timeout=10)
        try:
            await conn.fetchval("SELECT 1")
            ver = await conn.fetchval("SELECT version()")
        finally:
            await conn.close()
        return ver or ""

    try:
        ver = asyncio.run(asyncio.wait_for(_run(), timeout=15))
        short = str(ver).split(" on ")[0] if ver else ""
        _out(True, f"Connected to PostgreSQL{(' - ' + short) if short else ''}")
    except Exception as e:  # noqa: BLE001
        _out(False, f"PostgreSQL connection failed: {_friendly_conn_error(e, host, port, 'PostgreSQL')}")


# ── MySQL / MariaDB (PyMySQL - pure-python, sync) ────────────────────────────
def probe_mysql(creds):
    dsn = _pick(creds, "dsn", "MYSQL_DSN")
    if not dsn:
        _out(False, "Connection DSN required")
    try:
        import pymysql  # type: ignore
    except ImportError:
        _out(False, "PyMySQL not installed on this runner (pip install PyMySQL)")
    p = urllib.parse.urlparse(dsn)
    host, port = (p.hostname or ""), (p.port or 3306)
    if not host:
        _out(False, "MySQL connection failed: no host in the DSN - expected mysql://user:pass@HOST:3306/db")
    try:
        conn = pymysql.connect(
            host=host,
            port=port,
            user=urllib.parse.unquote(p.username or ""),
            password=urllib.parse.unquote(p.password or ""),
            db=(p.path or "").lstrip("/") or None,
            connect_timeout=10,
            read_timeout=10,
        )
        try:
            with conn.cursor() as cur:
                cur.execute("SELECT VERSION()")
                row = cur.fetchone()
            ver = (row[0] if row else "") or ""
        finally:
            conn.close()
        _out(True, f"Connected to MySQL{(' - ' + str(ver)) if ver else ''}")
    except Exception as e:  # noqa: BLE001
        _out(False, f"MySQL connection failed: {_friendly_conn_error(e, host, port, 'MySQL')}")


# ── Snowflake (key-pair JWT, mirrors the server handler) ─────────────────────
def probe_snowflake(creds):
    account = _pick(creds, "account", "SNOWFLAKE_ACCOUNT")
    user = _pick(creds, "user", "SNOWFLAKE_USER")
    pk = _pick(creds, "private_key", "SNOWFLAKE_PRIVATE_KEY")
    if not account or not user or not pk:
        _out(False, "Account + Username + Private Key required")
    try:
        from cryptography.hazmat.primitives import serialization  # type: ignore
        import requests  # type: ignore
    except ImportError:
        _out(False, "cryptography/requests not installed on this runner")

    def b64url(b):
        if isinstance(b, str):
            b = b.encode()
        return base64.urlsafe_b64encode(b).rstrip(b"=").decode()

    try:
        import time
        key = serialization.load_pem_private_key(pk.encode(), password=None)
        der_spki = key.public_key().public_bytes(
            serialization.Encoding.DER,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        fp = "SHA256:" + base64.b64encode(hashlib.sha256(der_spki).digest()).decode()
        acct_upper = account.upper()
        if "." in acct_upper:
            acct_upper = acct_upper.split(".")[0]
        qual = f"{acct_upper}.{user.upper()}"
        now = int(time.time())
        header = b64url(json.dumps({"alg": "RS256", "typ": "JWT"}))
        payload = b64url(json.dumps(
            {"iss": f"{qual}.{fp}", "sub": qual, "iat": now, "exp": now + 3540}))
        signing_input = f"{header}.{payload}".encode()
        from cryptography.hazmat.primitives import hashes  # type: ignore
        from cryptography.hazmat.primitives.asymmetric import padding  # type: ignore
        sig = key.sign(signing_input, padding.PKCS1v15(), hashes.SHA256())
        jwt = f"{header}.{payload}.{b64url(sig)}"
        host = f"https://{account.replace('_', '-').lower()}.snowflakecomputing.com"
        r = requests.post(
            f"{host}/api/v2/statements",
            headers={
                "Authorization": f"Bearer {jwt}",
                "X-Snowflake-Authorization-Token-Type": "KEYPAIR_JWT",
                "Content-Type": "application/json",
                "Accept": "application/json",
            },
            json={"statement": "SELECT CURRENT_VERSION()", "timeout": 20},
            timeout=15,
        )
        if r.ok:
            v = ""
            try:
                v = (r.json().get("data") or [[None]])[0][0] or ""
            except Exception:  # noqa: BLE001
                pass
            _out(True, f"Connected to Snowflake{(' (v' + str(v) + ')') if v else ''}")
        _out(False, f"Snowflake auth/query failed: HTTP {r.status_code} {r.text[:160]}")
    except Exception as e:  # noqa: BLE001
        _out(False, f"Snowflake connection failed: {e}")


# ── Databricks (PAT or OAuth M2M, mirrors the server handler) ────────────────
def probe_databricks(creds):
    import re
    from urllib.parse import urlsplit
    host = _pick(creds, "workspace_url", "DATABRICKS_HOST").rstrip("/")
    pat = _pick(creds, "token", "DATABRICKS_TOKEN")
    cid = _pick(creds, "client_id", "DATABRICKS_CLIENT_ID")
    sec = _pick(creds, "client_secret", "DATABRICKS_CLIENT_SECRET")
    if not host:
        _out(False, "Workspace URL required")
    url = urlsplit(host)
    if (url.scheme != "https" or url.username or url.password or url.port or url.path or url.query or url.fragment
        or not re.fullmatch(r"(?!accounts(?:\.|-))[a-z0-9-]+(?:\.[a-z0-9-]+)*\.(?:cloud\.databricks\.com|azuredatabricks\.net|gcp\.databricks\.com)", url.hostname or "")):
        _out(False, "Use an HTTPS Databricks workspace URL without a path or query")
    telemetry = {"User-Agent": "MelayaLabs_Melaya/2.0.0"}
    try:
        import requests  # type: ignore
    except ImportError:
        _out(False, "requests not installed on this runner")
    try:
        bearer = pat
        if not bearer:
            if not cid or not sec:
                _out(False, "Provide a Personal Access Token, or Client ID + Client Secret")
            tr = requests.post(
                f"{host}/oidc/v1/token",
                auth=(cid, sec),
                data={"grant_type": "client_credentials", "scope": "all-apis"},
                headers=telemetry, allow_redirects=False,
                timeout=12,
            )
            tj = {}
            try:
                tj = tr.json()
            except Exception:  # noqa: BLE001
                pass
            if tr.status_code != 200 or not tj.get("access_token"):
                _out(False, f"Databricks OAuth failed: HTTP {tr.status_code}")
            bearer = tj["access_token"]
        r = requests.get(
            f"{host}/api/2.1/clusters/list?page_size=1",
            headers={"Authorization": f"Bearer {bearer}", **telemetry}, allow_redirects=False,
            timeout=12,
        )
        if r.status_code == 200:
            _out(True, f"Connected to Databricks ({host.replace('https://', '').replace('http://', '')})")
        _out(False, f"Databricks unreachable: HTTP {r.status_code} {r.text[:120]}")
    except Exception as e:  # noqa: BLE001
        _out(False, f"Databricks connection failed: {e}")


# ── SQLite (local file on the runner) ────────────────────────────────────────
def probe_sqlite(creds):
    import os
    import sqlite3
    path = _pick(creds, "dsn", "db_path", "path")
    if not path:
        _out(False, "Database file path required")
    if not os.path.exists(path):
        _out(False, f"No SQLite file at {path}")
    try:
        conn = sqlite3.connect(path, timeout=8)
        try:
            conn.execute("SELECT 1")
        finally:
            conn.close()
        _out(True, f"Opened SQLite database ({os.path.basename(path)})")
    except Exception as e:  # noqa: BLE001
        _out(False, f"SQLite open failed: {e}")


_PROBES = {
    "postgres": probe_postgres,
    "mysql": probe_mysql,
    "snowflake": probe_snowflake,
    "databricks": probe_databricks,
    "sqlite": probe_sqlite,
}


def main():
    try:
        job = json.loads(sys.stdin.read() or "{}")
    except Exception as e:  # noqa: BLE001
        _out(False, f"bad probe job: {e}")
        return
    service = (job.get("service") or "").strip()
    creds = job.get("creds") or {}
    fn = _PROBES.get(service)
    if not fn:
        _out(False, f"unsupported service: {service}")
        return
    fn(creds)


if __name__ == "__main__":
    main()