backend/core/agent_tools.py

"""
Tools the geospatial agent can call.

Each tool pairs a Gemini `FunctionDeclaration` (what the model sees) with a Python
implementation (what actually runs). The agent decides which to call and reacts to
what comes back, instead of following a fixed detect-intent -> pick-tables ->
write-SQL pipeline.

The point of the toolset is that the agent can **look before it commits**. Every
bug class the old pipeline hit came from writing SQL blind against a schema
summary: guessing `US-%` when the data uses `USA-%`, colouring by `week_index`
because it happened to be the first numeric column, assuming a full-year raster
exists for a resident species. `sample_values` and `describe_table` make those
answerable in one cheap call rather than a wrong answer delivered confidently.

Results are deliberately compact. `run_sql` returns a preview and a row count,
never the full payload — the geometry goes to the map via `add_map_layer`, and
stuffing thousands of rows back into the prompt would blow the context and teach
the model nothing extra.
"""

from __future__ import annotations

import asyncio
import json
import logging
import re
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

from google.genai import types

from backend.core.jsonutil import dumps_safe, json_safe

logger = logging.getLogger(__name__)

# A tool call that runs longer than this is almost certainly a runaway spatial
# join; failing it returns control to the agent, which can narrow the query.
TOOL_TIMEOUT_SECONDS = 90.0

# Rows echoed back into the conversation from run_sql. Enough to see the shape of
# the result and spot obvious mistakes, small enough not to crowd the context.
PREVIEW_ROWS = 8

# Statements the agent may execute. Anything that writes is rejected outright:
# the agent is an analyst, not an administrator.
_FORBIDDEN_SQL = re.compile(
    r"\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|ATTACH|COPY|INSTALL|LOAD|PRAGMA|EXPORT)\b",
    re.IGNORECASE,
)


class ToolError(Exception):
    """A tool failed in a way the agent can reason about and retry."""


@dataclass
class Tool:
    """One callable tool: its schema for the model, and its implementation."""
    name: str
    description: str
    parameters: Dict[str, Any]
    run: Callable[..., Any]
    # Tools that produce user-visible output rather than just information.
    produces_output: bool = False

    def declaration(self) -> types.FunctionDeclaration:
        return types.FunctionDeclaration(
            name=self.name,
            description=self.description,
            parameters=self.parameters,
        )


@dataclass
class AgentContext:
    """
    Mutable state shared across one agent run.

    Collects everything the tools produce so the caller can emit a single final
    response: the map layer, chart, citations and stats accumulate here rather
    than being returned through the model, which would mean serializing large
    payloads into the prompt.
    """
    allowed_datasets: Optional[List[str]] = None
    # A list, not a single slot: the agent legitimately builds several layers for
    # one answer (one per species in a comparison). Keeping only the newest threw
    # away work it had already done and made the map contradict the text.
    layers: List[Dict[str, Any]] = field(default_factory=list)
    chart_data: Optional[Dict[str, Any]] = None
    raw_data: List[Dict[str, Any]] = field(default_factory=list)
    sql_statements: List[str] = field(default_factory=list)
    tables_used: List[str] = field(default_factory=list)
    stats: Dict[str, Any] = field(default_factory=dict)
    pending_question: Optional[Dict[str, Any]] = None


def _validate_sql(sql: str) -> str:
    """Reject anything that is not a read-only query."""
    cleaned = sql.strip().rstrip(";").strip()
    if not cleaned:
        raise ToolError("Empty SQL.")
    if _FORBIDDEN_SQL.search(cleaned):
        raise ToolError(
            "Only read-only SELECT/WITH queries are permitted. "
            "Remove any statement that modifies data or schema."
        )
    first = cleaned.lstrip("(").split(None, 1)[0].upper()
    if first not in ("SELECT", "WITH"):
        raise ToolError(f"Query must start with SELECT or WITH, got '{first}'.")
    return cleaned


def build_tools(ctx: AgentContext) -> Dict[str, Tool]:
    """
    Construct the toolset bound to one agent run.

    Imports are local so this module can be imported without spinning up DuckDB
    or the embedding index (useful for tests and for the schema-only path).
    """
    from backend.core.data_catalog import get_data_catalog, describe_table
    from backend.core.geo_engine import get_geo_engine
    from backend.core.semantic_search import get_semantic_search
    from backend.services.response_formatter import ResponseFormatter

    catalog = get_data_catalog()
    engine = get_geo_engine()
    semantic = get_semantic_search()

    def _in_scope(name: str) -> bool:
        return ctx.allowed_datasets is None or name in ctx.allowed_datasets

    # ---------------------------------------------------------------- discovery

    def search_datasets(query: str, limit: int = 8) -> Dict[str, Any]:
        hits = semantic.search(query, top_k=max(1, min(limit, 20)),
                               allowed_datasets=ctx.allowed_datasets)
        results = []
        for name, score in hits:
            meta = catalog.get_table_metadata(name)
            if not meta:
                continue
            results.append({
                "table": name,
                "relevance": round(float(score), 3),
                "rows": meta.get("row_count"),
                "description": describe_table(meta)[:400],
            })
        if not results:
            return {"results": [], "hint": "Nothing matched. Try broader wording, "
                                           "or call describe_table on a known table."}
        return {"results": results}

    def describe_table_tool(table: str) -> Dict[str, Any]:
        meta = catalog.get_table_metadata(table)
        if not meta:
            close = [n for n in catalog.catalog if table.lower() in n.lower()][:5]
            raise ToolError(
                f"No table named '{table}'." +
                (f" Did you mean: {', '.join(close)}?" if close else
                 " Use search_datasets to find one.")
            )
        if not _in_scope(table):
            raise ToolError(f"'{table}' is outside the datasets selected for this session.")
        if not engine.ensure_table_loaded(table):
            raise ToolError(f"'{table}' is in the catalog but its data file could not be loaded.")

        columns = engine.describe_columns(table)
        has_geometry = any(c[0] in engine.GEOMETRY_COLUMNS for c in columns)
        return {
            "table": table,
            "rows": meta.get("row_count"),
            "spatial": has_geometry,
            "columns": [{"name": c[0], "type": c[1]} for c in columns],
            "description": describe_table(meta),
            "attribution": meta.get("attribution"),
        }

    def sample_values(table: str, column: str, limit: int = 15) -> Dict[str, Any]:
        """Distinct values of a column — how the agent learns real formats."""
        if not _in_scope(table):
            raise ToolError(f"'{table}' is outside the datasets selected for this session.")
        if not engine.ensure_table_loaded(table):
            raise ToolError(f"Could not load '{table}'.")
        limit = max(1, min(limit, 50))
        try:
            distinct = engine.fetch_all(
                f'SELECT DISTINCT "{column}" FROM "{table}" '
                f'WHERE "{column}" IS NOT NULL LIMIT {limit}'
            )
            total = engine.fetch_one(
                f'SELECT COUNT(DISTINCT "{column}") FROM "{table}"'
            )[0]
        except Exception as e:
            raise ToolError(f"Could not sample '{column}' from '{table}': {e}")

        values = [r[0] for r in distinct]
        out: Dict[str, Any] = {
            "table": table,
            "column": column,
            "distinct_count": total,
            "sample": [str(v) for v in values],
        }
        # Numeric columns: the range matters more than example values.
        if values and isinstance(values[0], (int, float)) and not isinstance(values[0], bool):
            lo, hi = engine.fetch_one(
                f'SELECT MIN("{column}"), MAX("{column}") FROM "{table}"'
            )
            out["min"], out["max"] = lo, hi
        if total > limit:
            out["note"] = f"Showing {len(values)} of {total} distinct values."
        return out

    # ----------------------------------------------------------------- querying

    def _execute(sql: str) -> Dict[str, Any]:
        from backend.core.geo_engine import ResultTooLargeError  # noqa: F401

        cleaned = _validate_sql(sql)
        for name in catalog.catalog:
            if re.search(rf'(?<![\w."]){re.escape(name)}(?![\w."])', cleaned, re.IGNORECASE):
                if not _in_scope(name):
                    raise ToolError(f"'{name}' is outside the datasets selected for this session.")
                engine.ensure_table_loaded(name)
        try:
            return engine.execute_spatial_query(cleaned)
        except ResultTooLargeError as e:
            # Not a failure of the query's logic — it is simply too coarse. The
            # message names the fix, so the agent rewrites rather than gives up.
            raise ToolError(str(e))
        except Exception as e:
            # Hand the database's own message back; it is usually enough for the
            # agent to fix the query itself.
            raise ToolError(f"SQL failed: {e}")

    def run_sql(sql: str, purpose: str = "") -> Dict[str, Any]:
        result = _execute(sql)
        features = result.get("features", [])
        props = result.get("properties", {}) or {}
        ctx.sql_statements.append(sql.strip())

        preview = [
            {k: v for k, v in (f.get("properties") or {}).items()}
            for f in features[:PREVIEW_ROWS]
        ]
        out: Dict[str, Any] = {
            "row_count": len(features),
            "has_geometry": any(f.get("geometry") for f in features),
            "columns": list(preview[0].keys()) if preview else [],
            "preview": preview,
        }
        if not features:
            out["hint"] = ("Zero rows. Check literal values with sample_values before "
                           "assuming a format, and verify filters match real data.")
        return out

    def add_map_layer(sql: str, name: str, color_by: str = "") -> Dict[str, Any]:
        result = _execute(sql)
        features = result.get("features", [])
        if not features:
            raise ToolError("Query returned no rows, so there is nothing to map.")
        if not any(f.get("geometry") for f in features):
            raise ToolError(
                "Result has no geometry and cannot be mapped. Either select a geometry "
                "column, or join to a boundary table to borrow one."
            )

        ctx.sql_statements.append(sql.strip())
        geojson, _layer_id, layer_name = ResponseFormatter.format_geojson_layer(
            name, result, features, name, "📍", None, color_by=color_by or None,
        )

        # Credit the tables THIS query read, not every table the agent touched
        # while answering. A three-species comparison runs one query per species,
        # and attributing the union to all of them told the reader a layer was
        # built from three datasets when it came from one — and made the layer's
        # species look ambiguous when it was not.
        geojson.setdefault("properties", {})["source_tables"] = (
            ResponseFormatter._tables_referenced_in_sql(sql, list(catalog.catalog.keys()))
        )
        ctx.layers.append(geojson)
        return {
            "layer": layer_name,
            "features": len(features),
            "layers_on_map": len(ctx.layers),
            "note": "Layer added to the map. Call again to add another for a "
                    "different species or season; existing layers are kept.",
        }

    def make_chart(sql: str, chart_type: str, title: str) -> Dict[str, Any]:
        result = _execute(sql)
        features = result.get("features", [])
        if not features:
            raise ToolError("Query returned no rows, so there is nothing to chart.")
        ctx.sql_statements.append(sql.strip())
        chart = ResponseFormatter.generate_chart_data(
            sql, features, title,
            {"use_chart": True, "type": chart_type, "title": title},
        )
        if not chart:
            raise ToolError(
                "Could not build a chart from that result — it needs a label column "
                "and a numeric column."
            )
        ctx.chart_data = chart
        ctx.raw_data = ResponseFormatter.prepare_raw_data(features)
        return {"chart": chart.get("type"), "title": chart.get("title"),
                "points": len(chart.get("data") or [])}

    def compute_stats(sql: str, column: str) -> Dict[str, Any]:
        result = _execute(sql)
        features = result.get("features", [])
        values = [
            f["properties"].get(column) for f in features
            if isinstance(f.get("properties", {}).get(column), (int, float))
            and not isinstance(f["properties"].get(column), bool)
        ]
        if not values:
            raise ToolError(f"No numeric values found in column '{column}'.")
        ctx.sql_statements.append(sql.strip())
        values.sort()
        n = len(values)
        stats = {
            "count": n,
            "min": values[0],
            "max": values[-1],
            "mean": sum(values) / n,
            "median": values[n // 2] if n % 2 else (values[n // 2 - 1] + values[n // 2]) / 2,
        }
        ctx.stats[column] = stats
        return stats

    def ask_user(question: str, options: Optional[List[str]] = None) -> Dict[str, Any]:
        """Record a clarifying question; the loop stops and surfaces it."""
        ctx.pending_question = {"question": question, "options": options or []}
        return {"asked": question, "note": "Stop and wait for the user's answer."}

    # ------------------------------------------------------------------ schemas

    def _obj(props: Dict[str, Any], required: List[str]) -> Dict[str, Any]:
        return {"type": "object", "properties": props, "required": required}

    _str = {"type": "string"}
    _int = {"type": "integer"}

    tools = [
        Tool(
            name="search_datasets",
            description=(
                "Find datasets relevant to a question by meaning. Use this first when you "
                "do not already know which table holds the answer."
            ),
            parameters=_obj({
                "query": {**_str, "description": "What you are looking for, in plain language."},
                "limit": {**_int, "description": "Max results (default 8)."},
            }, ["query"]),
            run=search_datasets,
        ),
        Tool(
            name="describe_table",
            description=(
                "Get a table's exact columns, types, row count, whether it has geometry, "
                "and its documented meaning. Call before writing SQL against an unfamiliar table."
            ),
            parameters=_obj({"table": _str}, ["table"]),
            run=describe_table_tool,
        ),
        Tool(
            name="sample_values",
            description=(
                "List real distinct values of a column (and min/max if numeric). Use this "
                "BEFORE filtering on a value whose exact format you are unsure of — codes, "
                "category names, dates. Guessing a format returns zero rows with no error."
            ),
            parameters=_obj({
                "table": _str,
                "column": _str,
                "limit": {**_int, "description": "Max distinct values (default 15)."},
            }, ["table", "column"]),
            run=sample_values,
        ),
        Tool(
            name="run_sql",
            description=(
                "Execute a read-only DuckDB SQL query and get back the row count plus a small "
                "preview. Use it to check an approach or compute an answer. It does NOT put "
                "anything on the map — use add_map_layer for that."
            ),
            parameters=_obj({
                "sql": {**_str, "description": "A SELECT or WITH query."},
                "purpose": {**_str, "description": "One short line on what this is for."},
            }, ["sql"]),
            run=run_sql,
        ),
        Tool(
            name="add_map_layer",
            description=(
                "Run a query and add its result to the map as a styled layer. The query MUST "
                "select a geometry column. For a temporal animation, return every time step "
                "with its date/step column rather than filtering to one."
            ),
            parameters=_obj({
                "sql": {**_str, "description": "A SELECT that includes a geometry column."},
                "name": {**_str, "description": "Short layer name, 1-4 words."},
                "color_by": {**_str, "description": "Column to colour by (optional)."},
            }, ["sql", "name"]),
            run=add_map_layer,
            produces_output=True,
        ),
        Tool(
            name="make_chart",
            description=(
                "Run a query and turn the result into a chart for the Plots tab. Use for "
                "rankings, comparisons and distributions."
            ),
            parameters=_obj({
                "sql": _str,
                "chart_type": {**_str, "enum": ["bar", "line", "pie", "histogram"]},
                "title": _str,
            }, ["sql", "chart_type", "title"]),
            run=make_chart,
            produces_output=True,
        ),
        Tool(
            name="compute_stats",
            description="Summary statistics (count, min, max, mean, median) for a numeric column.",
            parameters=_obj({"sql": _str, "column": _str}, ["sql", "column"]),
            run=compute_stats,
        ),
        Tool(
            name="ask_user",
            description=(
                "Ask ONE clarifying question when the request is genuinely ambiguous and "
                "guessing would waste the user's time. Prefer making a reasonable choice and "
                "saying what you chose."
            ),
            parameters=_obj({
                "question": _str,
                "options": {"type": "array", "items": _str},
            }, ["question"]),
            run=ask_user,
        ),
    ]
    return {t.name: t for t in tools}


async def call_tool(tool: Tool, args: Dict[str, Any]) -> Dict[str, Any]:
    """
    Run a tool off the event loop, returning either its result or a structured
    error. Errors are returned rather than raised so the agent can read what went
    wrong and try something else — an exception here would end the whole turn.
    """
    try:
        result = await asyncio.wait_for(
            asyncio.to_thread(tool.run, **args), timeout=TOOL_TIMEOUT_SECONDS
        )
        return {"ok": True, "result": result}
    except asyncio.TimeoutError:
        return {"ok": False, "error":
                f"{tool.name} timed out after {TOOL_TIMEOUT_SECONDS:.0f}s. "
                "Narrow the query (filter by region, season or species) and retry."}
    except ToolError as e:
        return {"ok": False, "error": str(e)}
    except TypeError as e:
        return {"ok": False, "error": f"Bad arguments for {tool.name}: {e}"}
    except Exception as e:  # noqa: BLE001 - surface anything to the agent
        logger.warning(f"Tool {tool.name} failed: {e}", exc_info=True)
        return {"ok": False, "error": f"{tool.name} failed: {e}"}


def serialize_result(payload: Dict[str, Any]) -> Dict[str, Any]:
    """
    Make a tool result safe to hand back to the model.

    Note the previous version checked with a bare `json.dumps`, which does NOT
    raise on NaN — it emits a bare `NaN` token that the API then rejects with
    `400 INVALID_ARGUMENT`. json_safe strips those first.
    """
    safe = json_safe(payload)
    try:
        json.dumps(safe, allow_nan=False)
        return safe
    except (TypeError, ValueError):
        return json.loads(dumps_safe(safe))