tools.py

"""Tool function definitions and the central TOOLS / TOOL_HANDLER registry.

Every tool callable by the agent is defined here (or imported from a domain
module).  The ``ToolRegistry`` class manages the OpenAI-format definitions and
the name→callable mapping.
"""

import ast
import json
import subprocess
import glob as glob_mod
from typing import Callable

import cron_system
import task_system

from config import WORKDIR, get_logger
from wcode_errors import WcodeError, ToolError, CronError

logger = get_logger(__name__)
from hooks import check_deny_list
from messaging import run_spawn_subagent as _module_spawn_subagent
from skills import run_load_skill as _module_load_skill
from utils import safe_path, tool_def


# ═══════════════════════════════════════════════════════════════════════════
# ToolRegistry class
# ═══════════════════════════════════════════════════════════════════════════


class ToolRegistry:
    """Central registry of OpenAI-format tool definitions and callable handlers.

    Accepts optional manager references for tools that depend on other
    subsystems (skills, cron, subagents).  When omitted, the tool falls
    back to module-level shims.

    Tools are registered via the :meth:`register` decorator::

        tools = ToolRegistry()

        @tools.register
        def run_bash(command: str) -> str: ...

    Usage::

        registry = ToolRegistry(skill_registry=skills, cron_scheduler=cron)
        registry.ensure_mcp_tools(mcp_manager)
        # registry.tools  → list of OpenAI tool defs
        # registry.execute(tool_call_block) → tool result string
    """

    _class_functions: list[Callable] = []

    def __init__(
        self,
        skill_registry=None,
        cron_scheduler=None,
        subagent_manager=None,
    ) -> None:
        self.tools: list[dict] = []
        self.handlers: dict[str, Callable] = {}
        self._mcp_registered = False

        # ── Register all @tools.register decorated functions ──
        for fn in self._class_functions:
            self._register_fn(fn)

        # ── Apply manager-specific handler overrides ──
        if skill_registry is not None:
            self.handlers["run_load_skill"] = skill_registry.load_skill

        if cron_scheduler is not None:
            self._override_cron_handlers(cron_scheduler)

        if subagent_manager is not None:
            self.handlers["run_spawn_subagent"] = subagent_manager.spawn

    # ── Registration ───────────────────────────────────────────────────

    def _register_fn(self, fn: Callable) -> None:
        """Add a function to both the tool definitions and handler map."""
        name = fn.__name__
        self.tools.append(tool_def(fn))
        if name not in self.handlers:
            self.handlers[name] = fn

    def register(self, func: Callable) -> Callable:
        """Decorator: register a tool function and return it unchanged.

        Usage::

            @tools.register
            def run_bash(command: str) -> str: ...
        """
        ToolRegistry._class_functions.append(func)
        self._register_fn(func)
        return func

    def _override_cron_handlers(self, cron) -> None:
        """Override cron handlers with manager-aware wrappers.

        Tool definitions are already registered via the decorator;
        this only swaps in handlers that delegate to the cron scheduler.
        """
        def _schedule(cron_expr: str, prompt: str,
                      recurring: bool = True, durable: bool = True) -> str:
            try:
                result = cron.schedule(cron_expr, prompt, recurring, durable)
            except CronError as e:
                return f"Error: {e}"
            return f"Scheduled {result.id}: '{cron_expr}' → {prompt}"

        def _list_crons() -> str:
            jobs = cron.list_jobs()
            if not jobs:
                return "No cron jobs. Use schedule_cron to add one."
            lines = []
            for j in jobs:
                tag = "recurring" if j.recurring else "one-shot"
                dur = "durable" if j.durable else "session"
                lines.append(
                    f"  {j.id}: '{j.cron}' → {j.prompt[:40]} [{tag}, {dur}]"
                )
            return "\n".join(lines)

        self.handlers["run_schedule_cron"] = _schedule
        self.handlers["run_list_crons"] = _list_crons
        self.handlers["run_cancel_cron"] = cron.cancel

    # ── MCP integration ────────────────────────────────────────────────

    def ensure_mcp_tools(self, mcp_manager) -> None:
        """Register MCP-backed tools (idempotent — safe to call multiple times)."""
        if self._mcp_registered:
            return
        try:
            from mcp_manager import register_mcp_tools_to
            register_mcp_tools_to(self.handlers, self.tools, mcp_manager)
            self._mcp_registered = True
        except Exception:
            logger.exception("[mcp] failed to register MCP tools")

    # ── Execution ──────────────────────────────────────────────────────

    def execute(self, block) -> str:
        """Execute a single tool call by name, returning the result string.

        All :class:`WcodeError` exceptions raised by tool functions are
        caught here and converted to user-facing error strings.
        """
        name = block.function.name
        args = json.loads(block.function.arguments)
        handler = self.handlers.get(name)
        if handler is None:
            return f"Error: Unknown tool '{name}'"
        try:
            return handler(**args)
        except WcodeError as e:
            logger.warning(f"[tool] {name} error: {e}")
            return f"Error: {e}"


# ═══════════════════════════════════════════════════════════════════════════
# Module-level ToolRegistry instance — decorate tool functions with this
# ═══════════════════════════════════════════════════════════════════════════

tools = ToolRegistry()

# Imported tool functions (defined in other modules)
tools.register(_module_spawn_subagent)
tools.register(_module_load_skill)


# ═══════════════════════════════════════════════════════════════════════════
# Tool functions (module-level, stateless)
# ═══════════════════════════════════════════════════════════════════════════


@tools.register
def run_bash(command: str) -> str:
    """Run a shell command inside the project workspace.

    Parameters
    ----------
    command : str
        The shell command to execute (e.g. ``ls``, ``pytest``, ``git status``).
        Runs with a 120-second timeout and output is capped at 50000 bytes.

    Returns
    -------
    str
        Combined stdout + stderr, or an error message if the command fails.
    """
    deny_msg = check_deny_list(command)
    if deny_msg:
        return deny_msg
    try:
        r = subprocess.run(
            command,
            shell=True,
            cwd=WORKDIR,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=120,
        )
        out = (r.stdout + r.stderr).strip()
        return out[:50000] if out else "(no output)"
    except subprocess.TimeoutExpired:
        raise ToolError(f"Timeout (120s) {command}")
    except (FileNotFoundError, OSError) as e:
        logger.exception(f"Command failed: {command}")
        raise ToolError(f"{e}, command: {command}")


@tools.register
def run_read_file(path: str, offset: int = None, limit: int = None) -> str:
    """Read the contents of a file inside the project workspace.

    Parameters
    ----------
    path : str
        Path to the file, relative to the workspace root.
    offset : int, optional
        1-based line number to start reading from. Defaults to the first line.
    limit : int, optional
        Maximum number of lines to return. Reads to end of file when omitted.
    """
    try:
        all_lines = safe_path(path).read_text().splitlines()
        start = (offset - 1) if offset and offset > 0 else 0
        if start >= len(all_lines):
            raise ToolError(
                f"offset {offset} exceeds file length ({len(all_lines)} lines)"
            )
        lines = all_lines[start:]
        if limit and limit < len(lines):
            lines = lines[:limit] + [
                f"... ({len(all_lines) - start - limit} more lines)"
            ]
        return "\n".join(lines)
    except ToolError:
        raise
    except Exception as e:
        logger.exception(f"Failed to read file: {path}")
        raise ToolError(str(e))


@tools.register
def run_write_file(path: str, content: str) -> str:
    """Write (create or overwrite) a file inside the project workspace.

    Parameters
    ----------
    path : str
        Path to the file, relative to the workspace root.
        Parent directories are created if they do not exist.
    content : str
        The text content to write to the file.
    """
    try:
        file_path = safe_path(path)
        file_path.parent.mkdir(parents=True, exist_ok=True)
        file_path.write_text(content)
        return f"Wrote {len(content)} bytes to {path}"
    except ToolError:
        raise
    except Exception as e:
        logger.exception(f"Failed to write file: {path}")
        raise ToolError(str(e))


@tools.register
def run_edit_file(path: str, old_text: str, new_text: str) -> str:
    """Replace the first occurrence of *old_text* with *new_text* in a file.

    Parameters
    ----------
    path : str
        Path to the file, relative to the workspace root.
    old_text : str
        The exact text to find and replace.
    new_text : str
        The replacement text.
    """
    try:
        file_path = safe_path(path)
        text = file_path.read_text()
        if old_text not in text:
            raise ToolError(f"text not found in {path}")
        file_path.write_text(text.replace(old_text, new_text, 1))
        return f"Edited {path}"
    except ToolError:
        raise
    except Exception as e:
        logger.exception(f"Failed to edit file: {path}")
        raise ToolError(str(e))


@tools.register
def run_glob(pattern: str, recursive: bool = True) -> str:
    """Find files matching a glob pattern inside the project workspace.

    Parameters
    ----------
    pattern : str
        The glob pattern to match (e.g. ``**/*.py``, ``src/*.ts``).
    recursive : bool, optional
        Whether to search directories recursively. True by default.
    """
    try:
        results = set()
        for match in glob_mod.glob(pattern, root_dir=WORKDIR, recursive=recursive):
            if (WORKDIR / match).resolve().is_relative_to(WORKDIR):
                results.add(match)
        return "\n".join(results) if results else "(no matches)"
    except ToolError:
        raise
    except Exception as e:
        logger.exception(f"Glob failed: {pattern}")
        raise ToolError(f"[{type(e).__name__}] {e}")


# ── Cron wrappers ──


@tools.register
def run_schedule_cron(
    cron: str, prompt: str, recurring: bool = True, durable: bool = True
) -> str:
    """Schedule a cron job to inject messages into the agent loop.

    Parameters
    ----------
    cron : str
        5-field cron expression: ``minute hour day-of-month month day-of-week``
        (e.g. ``0 9 * * *``, ``*/5 * * * *``).
    prompt : str
        Message text injected into the conversation when the cron fires.
    recurring : bool, optional
        True for recurring jobs, False for one-shot. True by default.
    durable : bool, optional
        True to persist the job to disk across restarts. True by default.
    """
    try:
        result = cron_system.schedule_job(cron, prompt, recurring, durable)
    except CronError as e:
        return f"Error: {e}"
    return f"Scheduled {result.id}: '{cron}' → {prompt}"


@tools.register
def run_cancel_cron(job_id: str) -> str:
    """Cancel a scheduled cron job by its ID.

    Parameters
    ----------
    job_id : str
        The job ID returned by ``schedule_cron``.
    """
    return cron_system.cancel_job(job_id)


@tools.register
def run_list_crons() -> str:
    """List all currently registered cron jobs with their IDs and schedules."""
    jobs = cron_system.list_jobs()
    if not jobs:
        return "No cron jobs. Use schedule_cron to add one."
    lines = []
    for j in jobs:
        tag = "recurring" if j.recurring else "one-shot"
        dur = "durable" if j.durable else "session"
        lines.append(
            f"  {j.id}: '{j.cron}' → {j.prompt[:40]} " f"[{tag}, {dur}]"
        )
    return "\n".join(lines)


# ── Task wrappers ──


@tools.register
def run_create_task(subject: str, description: str = "", blockedBy: list = None) -> str:
    """Create a new task with optional dependency blocking.

    Parameters
    ----------
    subject : str
        A short title for the task.
    description : str, optional
        A longer description of what the task involves.
    blockedBy : list of str, optional
        List of task IDs that must be completed before this task can start.
    """
    task = task_system.create_task(subject, description, blockedBy)
    return json.dumps(task_system.asdict(task), indent=2)


@tools.register
def run_list_tasks() -> str:
    """List all tasks with their status, owner, and dependency information."""
    tasks = task_system.list_tasks()
    if not tasks:
        return "No tasks found."
    return json.dumps([task_system.asdict(t) for t in tasks], indent=2)


@tools.register
def run_get_task(task_id: str) -> str:
    """Get the full details of a specific task by ID.

    Parameters
    ----------
    task_id : str
        The task ID returned by ``create_task``.
    """
    return task_system.get_task(task_id)


@tools.register
def run_claim_task(task_id: str) -> str:
    """Claim a pending task, changing its status to in_progress.

    Parameters
    ----------
    task_id : str
        The task ID to claim.
    """
    return task_system.claim_task(task_id)


@tools.register
def run_complete_task(task_id: str) -> str:
    """Complete an in-progress task, unblocking tasks that depend on it.

    Parameters
    ----------
    task_id : str
        The task ID to mark as completed.
    """
    return task_system.complete_task(task_id)


def _normalize_todos(todos):
    if isinstance(todos, str):
        try:
            todos = json.loads(todos)
        except json.JSONDecodeError:
            try:
                todos = ast.literal_eval(todos)
            except (SyntaxError, ValueError):
                raise ToolError("todos must be a list or JSON array string")
    if not isinstance(todos, list):
        raise ToolError("todos must be a list")
    for i, todo in enumerate(todos):
        if not isinstance(todo, dict):
            raise ToolError(f"todos[{i}] must be an object")
        if "content" not in todo or "status" not in todo:
            raise ToolError(f"todos[{i}] missing 'content' or 'status'")
        if todo["status"] not in ("pending", "in_progress", "completed"):
            raise ToolError(f"todos[{i}] has invalid status '{todo['status']}'")
    return todos


@tools.register
def run_todo_write(todos: list) -> str:
    """Create and manage a task list for the current session.

    Each todo item must have a ``content`` string and a ``status`` value
    of ``"pending"``, ``"in_progress"``, or ``"completed"``.

    Parameters
    ----------
    todos : list[dict]
        A list of todo objects, each with keys:

        - **content** (str): Description of the task.
        - **status** (str): One of ``"pending"``, ``"in_progress"``, ``"completed"``.
    """
    todos = _normalize_todos(todos)
    logger.info(f"[todo] updated {len(todos)} item(s)")

    lines = ["Current Tasks"]
    for t in todos:
        icon = {"pending": " ", "in_progress": "▸", "completed": "✓"}[t["status"]]
        lines.append(f"  [{icon}] {t['content']}")
    logger.info("\n".join(lines))

    return f"Updated {len(todos)} todos"


# ═══════════════════════════════════════════════════════════════════════════
# Module-level defaults (backward compatibility)
# ═══════════════════════════════════════════════════════════════════════════

TOOLS: list[dict] = tools.tools
TOOL_HANDLER: dict[str, Callable] = tools.handlers


def _get_default() -> ToolRegistry:
    """Return the module-level ToolRegistry singleton."""
    return tools


def _init_module_defaults() -> None:
    """One-time initialization of module-level TOOLS / TOOL_HANDLER.

    Now a thin guard — the real initialization happens via decorators at
    import time.  Kept for backward compatibility with any external caller.
    """
    global TOOLS, TOOL_HANDLER
    if not TOOLS:
        TOOLS = tools.tools
        TOOL_HANDLER = tools.handlers


_init_module_defaults()