backend/tools/registry.py

"""
backend/tools/registry.py

All tools available to the Executor agent.
Each tool:
  - Has a JSON schema (for OpenAI function calling)
  - Is async
  - Returns structured output with metadata
  - Handles errors gracefully
"""
from __future__ import annotations
import asyncio
import json
import math
import re
import time
from datetime import datetime
from typing import Any

import httpx
from langchain_core.tools import tool


# ── Tool result wrapper ───────────────────────────────────────────────────────

def ok(data: Any, metadata: dict | None = None) -> dict:
    return {"status": "ok", "data": data, "metadata": metadata or {}, "timestamp": time.time()}

def err(message: str, details: str = "") -> dict:
    return {"status": "error", "error": message, "details": details, "timestamp": time.time()}


# ── Web search ────────────────────────────────────────────────────────────────

async def web_search(query: str, max_results: int = 5) -> dict:
    """Search via Serper (Google), falling back to DuckDuckGo if unavailable."""
    from ..core.config import get_settings
    settings = get_settings()

    # ── Primary: Serper (Google Search API) ──────────────────────────────────
    if settings.serper_api_key:
        try:
            async with httpx.AsyncClient(timeout=10) as c:
                r = await c.post(
                    "https://google.serper.dev/search",
                    headers={"X-API-KEY": settings.serper_api_key, "Content-Type": "application/json"},
                    json={"q": query, "num": int(max_results)},
                )
            if r.status_code == 200:
                data = r.json()
                items = data.get("organic", [])
                formatted = "\n\n".join(
                    f"**{item.get('title', '')}**\n{item.get('link', '')}\n{item.get('snippet', '')}"
                    for item in items
                )
                return ok({"query": query, "results": formatted[:3000], "count": len(items), "source": "serper"})
        except Exception:
            pass  # fall through to DuckDuckGo

    # ── Fallback: DuckDuckGo ──────────────────────────────────────────────────
    try:
        from duckduckgo_search import DDGS
        loop = asyncio.get_event_loop()

        def _search():
            with DDGS() as ddgs:
                return list(ddgs.text(query, max_results=int(max_results)))

        results = await loop.run_in_executor(None, _search)
        formatted = "\n\n".join(
            f"**{r.get('title', '')}**\n{r.get('href', '')}\n{r.get('body', '')}"
            for r in results
        )
        return ok({"query": query, "results": formatted[:3000], "count": len(results), "source": "duckduckgo"})
    except Exception as e:
        return err(f"Search failed: {e}", str(e))


# ── URL fetch ─────────────────────────────────────────────────────────────────

async def fetch_url(url: str, max_chars: int = 4000) -> dict:
    """Fetch content from a URL and extract clean text."""
    try:
        from bs4 import BeautifulSoup
        async with httpx.AsyncClient(timeout=15, follow_redirects=True) as client:
            resp = await client.get(url, headers={"User-Agent": "Mozilla/5.0"})
            resp.raise_for_status()
        soup = BeautifulSoup(resp.text, "html.parser")
        for tag in soup(["script", "style", "nav", "footer"]):
            tag.decompose()
        text = " ".join(soup.get_text(separator=" ").split())
        return ok({"url": url, "content": text[:max_chars], "status_code": resp.status_code})
    except Exception as e:
        return err(f"Fetch failed: {e}", str(e))


# ── Calculator ────────────────────────────────────────────────────────────────

def calculate(expression: str) -> dict:
    """
    Safely evaluate a mathematical expression.
    Supports: +, -, *, /, **, sqrt, log, sin, cos, etc.
    """
    # Whitelist safe operations only
    safe_names = {
        "abs": abs, "round": round, "min": min, "max": max,
        "sum": sum, "sqrt": math.sqrt, "log": math.log,
        "log10": math.log10, "exp": math.exp,
        "sin": math.sin, "cos": math.cos, "tan": math.tan,
        "pi": math.pi, "e": math.e, "pow": math.pow, "floor": math.floor,
        "ceil": math.ceil,
    }
    # Remove any potentially dangerous tokens
    cleaned = re.sub(r'[^0-9+\-*/().,\s\w]', '', expression)
    try:
        result = eval(cleaned, {"__builtins__": {}}, safe_names)
        return ok({"expression": expression, "result": result})
    except Exception as e:
        return err(f"Calculation failed: {e}", f"Expression: {expression}")


# ── In-memory file system ─────────────────────────────────────────────────────

_file_store: dict[str, str] = {}


def write_file(filename: str, content: str) -> dict:
    """Write content to an in-memory file store."""
    _file_store[filename] = content
    return ok({"filename": filename, "bytes": len(content), "files": list(_file_store.keys())})


def read_file(filename: str) -> dict:
    """Read content from in-memory file store."""
    if filename not in _file_store:
        return err(f"File not found: {filename}", f"Available: {list(_file_store.keys())}")
    return ok({"filename": filename, "content": _file_store[filename]})


def list_files() -> dict:
    """List all files in the in-memory store."""
    files = [{"name": k, "size": len(v)} for k, v in _file_store.items()]
    return ok({"files": files, "count": len(files)})


# ── Code runner ───────────────────────────────────────────────────────────────

def run_python(code: str) -> dict:
    """
    Execute Python code in a sandbox with common data-science libs available.
    Blocks truly dangerous ops (os, subprocess, socket, open).
    """
    import io, contextlib, builtins
    # Block dangerous modules only
    blocked = ["subprocess", "socket", "requests"]
    for b in blocked:
        if f"import {b}" in code:
            return err(f"Blocked: '{b}' not allowed in sandbox")
    if "open(" in code:
        return err("Blocked: file 'open()' not allowed in sandbox")

    # Allow standard builtins + safe data libs
    import math, json
    from datetime import datetime as _dt
    safe_globals = {
        "__builtins__": builtins,   # full builtins so import works
        "math": math,
        "json": json,
        "datetime": _dt,
    }
    # Try to inject optional libs if installed
    for lib in ["pandas", "numpy", "matplotlib"]:
        try:
            import importlib
            safe_globals[lib.split(".")[0]] = importlib.import_module(lib)
        except ImportError:
            pass

    stdout_capture = io.StringIO()
    try:
        with contextlib.redirect_stdout(stdout_capture):
            exec(code, safe_globals)
        output = stdout_capture.getvalue()
        return ok({"output": output[:2000] or "(no output)", "code": code[:500]})
    except Exception as e:
        return err(f"Execution error: {type(e).__name__}: {e}", code[:200])


# ── Current date/time ─────────────────────────────────────────────────────────

def get_datetime() -> dict:
    """Get current date, time, and timezone info."""
    now = datetime.utcnow()
    return ok({
        "utc": now.isoformat(),
        "date": now.strftime("%Y-%m-%d"),
        "time": now.strftime("%H:%M:%S"),
        "day_of_week": now.strftime("%A"),
    })


# ── Tool registry ─────────────────────────────────────────────────────────────

TOOL_SCHEMAS = [
    {
        "type": "function",
        "function": {
            "name": "web_search",
            "description": "Search the internet for current information. Use for facts, news, research.",
            "parameters": {
                "type": "object",
                "required": ["query"],
                "properties": {
                    "query": {"type": "string", "description": "Search query string"},
                    "max_results": {"type": "integer", "default": 5},
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "fetch_url",
            "description": "Fetch and read content from a specific URL.",
            "parameters": {
                "type": "object",
                "required": ["url"],
                "properties": {
                    "url": {"type": "string"},
                    "max_chars": {"type": "integer", "default": 4000},
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "Evaluate a mathematical expression. Supports arithmetic, sqrt, log, trig.",
            "parameters": {
                "type": "object",
                "required": ["expression"],
                "properties": {
                    "expression": {"type": "string", "description": "Math expression, e.g. 'sqrt(144) + 2**8'"},
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "write_file",
            "description": "Save content to a named file for later use.",
            "parameters": {
                "type": "object",
                "required": ["filename", "content"],
                "properties": {
                    "filename": {"type": "string"},
                    "content": {"type": "string"},
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Read previously saved file content.",
            "parameters": {
                "type": "object",
                "required": ["filename"],
                "properties": {
                    "filename": {"type": "string"},
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "run_python",
            "description": "Execute Python code and capture output. Use for data processing, calculations.",
            "parameters": {
                "type": "object",
                "required": ["code"],
                "properties": {
                    "code": {"type": "string", "description": "Python code to execute"},
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_datetime",
            "description": "Get current UTC date and time.",
            "parameters": {"type": "object", "properties": {}},
        },
    },
]


async def execute_tool(name: str, args: dict) -> dict:
    """Dispatch a tool call by name."""
    dispatch = {
        "web_search": lambda: web_search(**args),
        "fetch_url": lambda: fetch_url(**args),
        "calculate": lambda: calculate(**args),
        "write_file": lambda: write_file(**args),
        "read_file": lambda: read_file(**args),
        "run_python": lambda: run_python(**args),
        "get_datetime": lambda: get_datetime(),
    }
    fn = dispatch.get(name)
    if fn is None:
        return err(f"Unknown tool: {name}")
    try:
        result = fn()
        if asyncio.iscoroutine(result):
            return await result
        return result
    except Exception as e:
        return err(f"Tool '{name}' crashed: {e}", str(e))