"""MySQL / PostgreSQL / SQLite 异步访问（供 Agent 在用户提供连接信息时使用）。

PostgreSQL 上带 RETURNING 的 INSERT 等会走 execute 路径，不返回结果集；需要时用只读 SELECT 再查。
"""

from __future__ import annotations

import asyncio
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Sequence

WRITE_CONFIRM = "EXECUTE_WRITE"
_MAX_SQL_CHARS = 200_000
_DEFAULT_TIMEOUT = 60.0


@dataclass(frozen=True)
class DbConnectionSpec:
    engine: str  # sqlite | postgresql | mysql
    path: str | None = None
    host: str | None = None
    port: int | None = None
    user: str | None = None
    password: str | None = None
    database: str | None = None
    schema: str | None = None  # PG default schema / search_path hint
    sslmode: str | None = None


def connection_from_dict(d: dict[str, Any]) -> DbConnectionSpec:
    eng = str(d.get("engine") or "").strip().lower()
    if eng in ("pg", "postgres"):
        eng = "postgresql"
    if eng not in ("sqlite", "postgresql", "mysql"):
        raise ValueError(f"unsupported engine: {eng!r} (use sqlite|postgresql|mysql)")
    path = d.get("path")
    path_s = str(path).strip() if path else None
    port = d.get("port")
    pi = int(port) if port is not None else None
    return DbConnectionSpec(
        engine=eng,
        path=path_s,
        host=str(d["host"]).strip() if d.get("host") else None,
        port=pi,
        user=str(d["user"]).strip() if d.get("user") else None,
        password=str(d["password"]) if d.get("password") is not None else None,
        database=str(d["database"]).strip() if d.get("database") else None,
        schema=str(d["schema"]).strip() if d.get("schema") else None,
        sslmode=str(d["sslmode"]).strip() if d.get("sslmode") else None,
    )


def validate_connection(spec: DbConnectionSpec) -> None:
    if spec.engine == "sqlite":
        if not spec.path:
            raise ValueError("sqlite requires connection.path")
        return
    if not spec.host:
        raise ValueError(f"{spec.engine} requires connection.host")
    if not spec.user:
        raise ValueError(f"{spec.engine} requires connection.user")
    if spec.password is None:
        raise ValueError(f"{spec.engine} requires connection.password (may be empty string)")
    if not spec.database:
        raise ValueError(f"{spec.engine} requires connection.database")


def _safe_ident(name: str) -> str:
    s = name.strip()
    if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", s):
        raise ValueError(f"invalid SQL identifier: {name!r}")
    return s


def _strip_leading_sql_comments(sql: str) -> str:
    s = sql.strip()
    while True:
        if s.startswith("--"):
            nl = s.find("\n")
            if nl == -1:
                return ""
            s = s[nl + 1 :].lstrip()
            continue
        if s.startswith("/*"):
            end = s.find("*/", 2)
            if end == -1:
                return s
            s = s[end + 2 :].lstrip()
            continue
        break
    return s


def first_statement_only(sql: str) -> str:
    """只取第一条语句（分号分隔），降低多语句注入面。"""
    body = sql.strip()
    depth = 0
    in_single = False
    in_double = False
    i = 0
    while i < len(body):
        c = body[i]
        if in_single:
            if c == "'" and body[i - 1 : i + 2] != "''":
                in_single = False
            i += 1
            continue
        if in_double:
            if c == '"':
                in_double = False
            i += 1
            continue
        if c == "'":
            in_single = True
            i += 1
            continue
        if c == '"':
            in_double = True
            i += 1
            continue
        if c == "(":
            depth += 1
        elif c == ")" and depth > 0:
            depth -= 1
        elif c == ";" and depth == 0:
            return body[:i].strip()
        i += 1
    return body.strip()


def is_readonly_sql(sql: str) -> bool:
    """粗略判断只读：首条语句以 SELECT/WITH/SHOW/DESC/DESCRIBE/EXPLAIN 开头。"""
    one = first_statement_only(sql)
    head = _strip_leading_sql_comments(one)
    if not head:
        return False
    return bool(
        re.match(
            r"(select|with|show|describe|desc|explain)\b",
            head,
            re.IGNORECASE | re.DOTALL,
        )
    )


def _json_cell(v: Any) -> Any:
    if v is None:
        return None
    if isinstance(v, (bool, int, float, str)):
        return v
    if isinstance(v, bytes):
        hx = v[:64].hex()
        return f"<bytes len={len(v)} hex64={hx}>"
    return str(v)


def _rows_to_json(
    columns: list[str], rows: Sequence[Sequence[Any]], max_rows: int
) -> tuple[list[dict[str, Any]], bool]:
    out: list[dict[str, Any]] = []
    truncated = False
    for i, row in enumerate(rows):
        if i >= max_rows:
            truncated = True
            break
        out.append({columns[j]: _json_cell(row[j]) for j in range(len(columns))})
    return out, truncated


async def _run_sqlite(
    path: Path,
    sql: str,
    *,
    read_only: bool,
    max_rows: int,
) -> dict[str, Any]:
    import aiosqlite

    mode = "ro" if read_only else "rwc"
    uri = f"file:{path.as_posix()}?mode={mode}"
    async with aiosqlite.connect(uri, uri=True) as db:
        db.row_factory = aiosqlite.Row
        try:
            async with db.execute(sql) as cur:
                if cur.description:
                    cols = [d[0] for d in cur.description]
                    raw = await cur.fetchmany(max_rows + 1)
                    rows, trunc = _rows_to_json(cols, raw, max_rows)
                    return {
                        "ok": True,
                        "columns": cols,
                        "rows": rows,
                        "truncated": trunc,
                        "rowcount": cur.rowcount if cur.rowcount >= 0 else None,
                    }
                rc = cur.rowcount
            if not read_only:
                await db.commit()
            return {
                "ok": True,
                "columns": [],
                "rows": [],
                "truncated": False,
                "rowcount": rc,
            }
        except Exception as e:  # noqa: BLE001
            await db.rollback()
            return {"ok": False, "error": repr(e)}


def _pg_stmt_returns_rows(stmt: str) -> bool:
    s = _strip_leading_sql_comments(stmt).lstrip().lower()
    if re.match(
        r"(select|with|show|values|table|explain|describe)\b",
        s,
    ):
        return True
    return s.startswith("desc ") or s.startswith("desc\t")


def _pg_stmt_is_mutating(stmt: str) -> bool:
    s = _strip_leading_sql_comments(stmt).lstrip().lower()
    return bool(
        re.match(
            r"(insert|update|delete|create|drop|alter|truncate|grant|revoke)\b",
            s,
        )
    )


async def _run_postgresql(
    spec: DbConnectionSpec,
    sql: str,
    *,
    read_only: bool,
    max_rows: int,
) -> dict[str, Any]:
    try:
        import asyncpg
    except ImportError as e:
        return {
            "ok": False,
            "error": f"asyncpg not installed: {e}. pip install 'aicd[db]'",
        }
    port = spec.port or 5432
    kwargs: dict[str, Any] = {
        "host": spec.host,
        "port": port,
        "user": spec.user,
        "password": spec.password,
        "database": spec.database,
        "timeout": _DEFAULT_TIMEOUT,
    }
    if spec.sslmode:
        kwargs["ssl"] = spec.sslmode != "disable"
    conn = await asyncpg.connect(**kwargs)
    try:
        if spec.schema:
            sch = _safe_ident(spec.schema)
            await conn.execute(f"SET search_path TO {sch}")
        stmt = first_statement_only(sql)
        use_fetch = read_only or (
            not _pg_stmt_is_mutating(stmt) and _pg_stmt_returns_rows(stmt)
        )
        if use_fetch:
            rows = await conn.fetch(stmt)
            if not rows:
                return {
                    "ok": True,
                    "columns": [],
                    "rows": [],
                    "truncated": False,
                    "rowcount": 0,
                }
            cols = list(rows[0].keys())
            data = [tuple(r[c] for c in cols) for r in rows[: max_rows + 1]]
            out, trunc = _rows_to_json(cols, data, max_rows)
            return {
                "ok": True,
                "columns": cols,
                "rows": out,
                "truncated": trunc,
                "rowcount": len(rows),
            }
        status = await conn.execute(stmt)
        return {
            "ok": True,
            "columns": [],
            "rows": [],
            "truncated": False,
            "status": str(status),
            "rowcount": None,
        }
    except Exception as e:  # noqa: BLE001
        return {"ok": False, "error": repr(e)}
    finally:
        await conn.close()


async def _run_mysql(
    spec: DbConnectionSpec,
    sql: str,
    *,
    read_only: bool,
    max_rows: int,
) -> dict[str, Any]:
    try:
        import aiomysql
    except ImportError as e:
        return {
            "ok": False,
            "error": f"aiomysql not installed: {e}. pip install 'aicd[db]'",
        }
    port = spec.port or 3306
    conn = await aiomysql.connect(
        host=spec.host,
        port=port,
        user=spec.user,
        password=spec.password,
        db=spec.database,
        autocommit=True,
    )
    try:
        async with conn.cursor(aiomysql.DictCursor) as cur:
            await cur.execute(sql)
            if cur.description:
                cols = [d[0] for d in cur.description]
                raw = await cur.fetchmany(max_rows + 1)
                tuples = [tuple(row[c] for c in cols) for row in raw]
                out, trunc = _rows_to_json(cols, tuples, max_rows)
                return {
                    "ok": True,
                    "columns": cols,
                    "rows": out,
                    "truncated": trunc,
                    "rowcount": cur.rowcount,
                }
            return {
                "ok": True,
                "columns": [],
                "rows": [],
                "truncated": False,
                "rowcount": cur.rowcount,
            }
    except Exception as e:  # noqa: BLE001
        return {"ok": False, "error": repr(e)}
    finally:
        conn.close()


async def db_execute(
    spec: DbConnectionSpec,
    sql: str,
    *,
    read_only: bool,
    max_rows: int,
    write_confirm: str | None,
) -> dict[str, Any]:
    if len(sql) > _MAX_SQL_CHARS:
        return {"ok": False, "error": f"sql exceeds {_MAX_SQL_CHARS} chars"}
    if not read_only:
        if write_confirm != WRITE_CONFIRM:
            return {
                "ok": False,
                "error": f"writes require write_confirm={WRITE_CONFIRM!r}",
            }
    else:
        if not is_readonly_sql(sql):
            return {
                "ok": False,
                "error": "read_only mode: only SELECT/WITH/SHOW/DESCRIBE/DESC/EXPLAIN allowed",
            }

    validate_connection(spec)
    full = sql.strip()
    stmt = first_statement_only(sql)
    if full != stmt:
        suffix = full[len(stmt) :].lstrip()
        if suffix.startswith(";") and suffix[1:].strip():
            return {"ok": False, "error": "multiple statements are not allowed"}

    async def _go() -> dict[str, Any]:
        if spec.engine == "sqlite":
            p = Path(spec.path or "")
            return await _run_sqlite(p, stmt, read_only=read_only, max_rows=max_rows)
        if spec.engine == "postgresql":
            return await _run_postgresql(
                spec, stmt, read_only=read_only, max_rows=max_rows
            )
        return await _run_mysql(spec, stmt, read_only=read_only, max_rows=max_rows)

    return await asyncio.wait_for(_go(), timeout=_DEFAULT_TIMEOUT + 5.0)


async def db_ping(spec: DbConnectionSpec) -> dict[str, Any]:
    validate_connection(spec)
    if spec.engine == "sqlite":
        import aiosqlite

        p = Path(spec.path or "")
        uri = f"file:{p.as_posix()}?mode=ro"
        try:
            async with aiosqlite.connect(uri, uri=True) as db:
                async with db.execute("SELECT 1") as cur:
                    await cur.fetchone()
            return {"ok": True, "engine": "sqlite", "path": str(p.resolve())}
        except Exception as e:  # noqa: BLE001
            return {"ok": False, "error": repr(e)}
    if spec.engine == "postgresql":
        r = await _run_postgresql(
            spec, "SELECT version()", read_only=True, max_rows=1
        )
        if r.get("ok") and r.get("rows"):
            return {"ok": True, "engine": "postgresql", "version": r["rows"][0]}
        return r
    r = await _run_mysql(spec, "SELECT VERSION() AS v", read_only=True, max_rows=1)
    if r.get("ok") and r.get("rows"):
        return {"ok": True, "engine": "mysql", "version": r["rows"][0]}
    return r


async def db_list_tables(spec: DbConnectionSpec) -> dict[str, Any]:
    validate_connection(spec)
    if spec.engine == "sqlite":
        sql = (
            "SELECT name FROM sqlite_master "
            "WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name"
        )
        p = Path(spec.path or "")
        return await _run_sqlite(p, sql, read_only=True, max_rows=10_000)
    if spec.engine == "postgresql":
        if spec.schema:
            sch = _safe_ident(spec.schema)
            sql = f"""
                SELECT table_schema, table_name
                FROM information_schema.tables
                WHERE table_type = 'BASE TABLE'
                  AND table_schema NOT IN ('pg_catalog', 'information_schema')
                  AND table_schema = '{sch}'
                ORDER BY table_schema, table_name
            """
        else:
            sql = """
                SELECT table_schema, table_name
                FROM information_schema.tables
                WHERE table_type = 'BASE TABLE'
                  AND table_schema NOT IN ('pg_catalog', 'information_schema')
                ORDER BY table_schema, table_name
            """
        return await _run_postgresql(spec, sql, read_only=True, max_rows=10_000)
    sql = (
        "SELECT TABLE_SCHEMA AS table_schema, TABLE_NAME AS table_name "
        "FROM information_schema.tables "
        "WHERE table_type = 'BASE TABLE' AND TABLE_SCHEMA = DATABASE() "
        "ORDER BY TABLE_NAME"
    )
    return await _run_mysql(spec, sql, read_only=True, max_rows=10_000)


async def db_list_schemas(spec: DbConnectionSpec) -> dict[str, Any]:
    validate_connection(spec)
    if spec.engine == "sqlite":
        return {"ok": True, "schemas": [{"name": "main"}]}
    if spec.engine == "postgresql":
        sql = (
            "SELECT schema_name FROM information_schema.schemata "
            "WHERE schema_name NOT LIKE 'pg_%' AND schema_name != 'information_schema' "
            "ORDER BY schema_name"
        )
        return await _run_postgresql(spec, sql, read_only=True, max_rows=2000)
    sql = "SHOW DATABASES"
    return await _run_mysql(spec, sql, read_only=True, max_rows=2000)


async def db_describe_table(
    spec: DbConnectionSpec, table: str, schema: str | None = None
) -> dict[str, Any]:
    t = _safe_ident(table)
    validate_connection(spec)
    if spec.engine == "sqlite":
        sql = f"PRAGMA table_info({t})"
        p = Path(spec.path or "")
        return await _run_sqlite(p, sql, read_only=True, max_rows=500)
    if spec.engine == "postgresql":
        sch = _safe_ident(schema or spec.schema or "public")
        sql = f"""
            SELECT column_name, data_type, is_nullable, column_default
            FROM information_schema.columns
            WHERE table_schema = '{sch}' AND table_name = '{t}'
            ORDER BY ordinal_position
        """
        return await _run_postgresql(spec, sql, read_only=True, max_rows=500)
    sql = f"""
        SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_DEFAULT
        FROM information_schema.COLUMNS
        WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = '{t}'
        ORDER BY ORDINAL_POSITION
    """
    return await _run_mysql(spec, sql, read_only=True, max_rows=500)


async def db_row_count(
    spec: DbConnectionSpec, table: str, schema: str | None = None
) -> dict[str, Any]:
    t = _safe_ident(table)
    validate_connection(spec)
    if spec.engine == "sqlite":
        sql = f'SELECT COUNT(*) AS cnt FROM "{t}"'
        p = Path(spec.path or "")
        return await _run_sqlite(p, sql, read_only=True, max_rows=1)
    if spec.engine == "postgresql":
        sch = _safe_ident(schema or spec.schema or "public")
        sql = f'SELECT COUNT(*) AS cnt FROM "{sch}"."{t}"'
        return await _run_postgresql(spec, sql, read_only=True, max_rows=1)
    sql = f"SELECT COUNT(*) AS cnt FROM `{t}`"
    return await _run_mysql(spec, sql, read_only=True, max_rows=1)


async def db_inspect_operation(
    spec: DbConnectionSpec,
    operation: str,
    *,
    table: str | None = None,
    schema: str | None = None,
) -> dict[str, Any]:
    op = operation.strip().lower()
    if op == "ping":
        return await db_ping(spec)
    if op == "list_tables":
        return await db_list_tables(spec)
    if op == "list_schemas":
        return await db_list_schemas(spec)
    if op == "describe_table":
        if not table:
            return {"ok": False, "error": "describe_table requires table"}
        return await db_describe_table(spec, table, schema)
    if op in ("row_count", "table_row_count"):
        if not table:
            return {"ok": False, "error": "row_count requires table"}
        return await db_row_count(spec, table, schema)
    return {"ok": False, "error": f"unknown operation: {operation}"}
