from __future__ import annotations
from typing import Any, Callable, Dict, Mapping, Optional

from .client import AgenticBQ
from .models import QueryConfig

def bigquery_toolkit(
    bq: AgenticBQ,
    default_config: Optional[QueryConfig] = None,
) -> Dict[str, Callable[..., Any]]:
    cfg = default_config or QueryConfig()

    def tool_safe_query(sql: str, params: Optional[Mapping[str, Any]] = None) -> dict[str, Any]:
        res = bq.safe_query(sql=sql, params=params, config=cfg)
        return {
            "job_id": res.job_id,
            "total_bytes_processed": res.total_bytes_processed,
            "sql": res.sql,
            "rows": res.rows,
        }

    def tool_list_datasets() -> list[dict[str, Any]]:
        return [di.__dict__ for di in bq.list_datasets()]

    def tool_list_tables(dataset: str) -> list[dict[str, Any]]:
        return [ti.__dict__ for ti in bq.list_tables(dataset)]

    def tool_table_schema(dataset: str, table: str) -> dict[str, Any]:
        ts = bq.table_schema(dataset, table)
        return {
            "project": ts.project,
            "dataset": ts.dataset,
            "table": ts.table,
            "fields": [f.__dict__ for f in ts.fields],
        }

    return {
        "bq_safe_query": tool_safe_query,
        "bq_list_datasets": tool_list_datasets,
        "bq_list_tables": tool_list_tables,
        "bq_table_schema": tool_table_schema,
    }

