"""NATIVE handler: system.db_query — Execute SQL queries on databases.

Connects to a database accessible from the customer's network and runs
parameterized SQL queries.  This is a key qRemoteX use case: the hosted
qRunX cannot reach private databases, but qRemoteX (in the customer's VPC) can.

Architecture reference: Section 12.2 — Tool Categories & Execution Matrix.

Security:
  - Uses parameterized queries only (no string interpolation)
  - Connection strings come from pre-resolved tool_params (injected by qRunX)
  - No credentials are logged
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any
from urllib.parse import urlparse

logger = logging.getLogger("qremotex.tools.db_query")

# Supported database drivers (optional imports)
_DRIVER_AVAILABLE: dict[str, bool] = {}


def _check_drivers() -> None:
    """Check which database drivers are available at import time."""
    global _DRIVER_AVAILABLE
    for mod_name in ("asyncpg", "aiosqlite", "aiomysql"):
        try:
            __import__(mod_name)
            _DRIVER_AVAILABLE[mod_name] = True
        except ImportError:
            _DRIVER_AVAILABLE[mod_name] = False


_check_drivers()


async def handle_db_query(
    tool_params: dict[str, Any],
    context: dict[str, Any],
) -> dict[str, Any]:
    """Execute a database query.

    Expected tool_params:
        query (str): SQL query with $1/$2 or ? placeholders
        params (list): Bind parameters (optional, default [])
        connection_string (str): Database connection URI
            e.g., "postgresql://user:pass@10.0.1.50:5432/mydb"
        max_rows (int): Maximum rows to return (default 1000)

    Returns:
        dict with keys: rows (list[dict]), row_count (int), columns (list[str])
    """
    query = tool_params.get("query")
    if not query:
        return {"error": "Missing required parameter: query", "rows": [], "row_count": 0}

    connection_string = tool_params.get("connection_string", "")
    if not connection_string:
        return {"error": "Missing required parameter: connection_string", "rows": [], "row_count": 0}

    params = tool_params.get("params", [])
    max_rows = min(int(tool_params.get("max_rows", 1000)), 10000)

    # Parse connection string to determine driver
    parsed = urlparse(connection_string)
    scheme = parsed.scheme.lower()

    try:
        if scheme in ("postgresql", "postgres"):
            return await _query_postgres(connection_string, query, params, max_rows)
        elif scheme == "sqlite":
            return await _query_sqlite(parsed.path, query, params, max_rows)
        elif scheme == "mysql":
            return await _query_mysql(connection_string, query, params, max_rows)
        else:
            return {
                "error": f"Unsupported database scheme: {scheme}. Supported: postgresql, sqlite, mysql",
                "rows": [],
                "row_count": 0,
            }
    except Exception as e:
        logger.error("db_query failed: %s", type(e).__name__)
        return {
            "error": f"Database query failed: {type(e).__name__}: {e}",
            "rows": [],
            "row_count": 0,
        }


async def _query_postgres(
    dsn: str, query: str, params: list, max_rows: int,
) -> dict[str, Any]:
    """Execute query via asyncpg."""
    if not _DRIVER_AVAILABLE.get("asyncpg"):
        return {"error": "asyncpg driver not installed. Run: pip install asyncpg", "rows": [], "row_count": 0}

    import asyncpg

    conn = await asyncio.wait_for(asyncpg.connect(dsn), timeout=10)
    try:
        stmt = await conn.prepare(query)
        records = await stmt.fetch(*params)
        columns = [attr.name for attr in stmt.get_attributes()]

        rows = [dict(zip(columns, row)) for row in records[:max_rows]]
        # Convert non-JSON-serializable types to strings
        for row in rows:
            for k, v in row.items():
                if not isinstance(v, (str, int, float, bool, type(None), list, dict)):
                    row[k] = str(v)

        return {
            "rows": rows,
            "row_count": len(rows),
            "columns": columns,
            "total_rows_fetched": len(records),
            "truncated": len(records) > max_rows,
        }
    finally:
        await conn.close()


async def _query_sqlite(
    db_path: str, query: str, params: list, max_rows: int,
) -> dict[str, Any]:
    """Execute query via aiosqlite."""
    if not _DRIVER_AVAILABLE.get("aiosqlite"):
        return {"error": "aiosqlite driver not installed. Run: pip install aiosqlite", "rows": [], "row_count": 0}

    import aiosqlite

    async with aiosqlite.connect(db_path) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute(query, params)
        raw_rows = await cursor.fetchmany(max_rows)
        columns = [desc[0] for desc in cursor.description] if cursor.description else []

        rows = [dict(zip(columns, row)) for row in raw_rows]
        return {
            "rows": rows,
            "row_count": len(rows),
            "columns": columns,
            "truncated": False,
        }


async def _query_mysql(
    dsn: str, query: str, params: list, max_rows: int,
) -> dict[str, Any]:
    """Execute query via aiomysql."""
    if not _DRIVER_AVAILABLE.get("aiomysql"):
        return {"error": "aiomysql driver not installed. Run: pip install aiomysql", "rows": [], "row_count": 0}

    import aiomysql
    parsed = urlparse(dsn)

    conn = await asyncio.wait_for(
        aiomysql.connect(
            host=parsed.hostname or "localhost",
            port=parsed.port or 3306,
            user=parsed.username or "",
            password=parsed.password or "",
            db=parsed.path.lstrip("/") or "",
        ),
        timeout=10,
    )
    try:
        async with conn.cursor(aiomysql.DictCursor) as cur:
            await cur.execute(query, params or None)
            rows = await cur.fetchmany(max_rows)
            columns = [desc[0] for desc in cur.description] if cur.description else []
            return {
                "rows": list(rows),
                "row_count": len(rows),
                "columns": columns,
                "truncated": False,
            }
    finally:
        conn.close()
