from __future__ import annotations

from dataclasses import replace
from typing import TYPE_CHECKING, Any

from aicd.tools import db_multi

if TYPE_CHECKING:
    from aicd.agent.tool_router import ToolRouter


def _conn_dict(args: dict[str, Any]) -> dict[str, Any]:
    c = args.get("connection")
    if not isinstance(c, dict):
        raise ValueError("connection must be an object")
    return c


async def _db_inspect(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    try:
        raw = _conn_dict(args)
        spec = db_multi.connection_from_dict(raw)
        if spec.engine == "sqlite" and spec.path:
            p = router._resolve_path(spec.path)
            spec = replace(spec, path=str(p))
        op = str(args.get("operation") or "")
        return await db_multi.db_inspect_operation(
            spec,
            op,
            table=args.get("table"),
            schema=args.get("schema"),
        )
    except PermissionError as e:
        return {"ok": False, "error": str(e)}
    except ValueError as e:
        return {"ok": False, "error": str(e)}
    except Exception as e:  # noqa: BLE001
        return {"ok": False, "error": repr(e)}


async def _db_query(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    try:
        raw = _conn_dict(args)
        spec = db_multi.connection_from_dict(raw)
        if spec.engine == "sqlite" and spec.path:
            p = router._resolve_path(spec.path)
            spec = replace(spec, path=str(p))
        sql = str(args.get("sql") or "")
        read_only = bool(args.get("read_only", True))
        max_rows = int(args.get("max_rows", 500))
        max_rows = max(1, min(max_rows, 20_000))
        wc = args.get("write_confirm")
        wc_s = str(wc) if wc is not None else None
        return await db_multi.db_execute(
            spec,
            sql,
            read_only=read_only,
            max_rows=max_rows,
            write_confirm=wc_s,
        )
    except PermissionError as e:
        return {"ok": False, "error": str(e)}
    except ValueError as e:
        return {"ok": False, "error": str(e)}
    except Exception as e:  # noqa: BLE001
        return {"ok": False, "error": repr(e)}


HANDLERS: dict[str, Any] = {
    "db_inspect": _db_inspect,
    "db_query": _db_query,
}
