"""Registry mapping step kinds to handler functions."""
from __future__ import annotations

from typing import Callable, Optional

from core.state import PlanStep, StepResult
from runtime.tools.file_ops import finish_step, read_file, write_file
from runtime.tools.flaky import run_flaky
from runtime.tools.shell import run_shell
from runtime.tools.web_fetch import web_fetch
from runtime.tools.web_search import web_search

ToolHandler = Callable[[PlanStep], StepResult]


class ToolRegistry:
    """Maps step kind strings to callable tool handlers."""

    def __init__(self) -> None:
        self._handlers: dict[str, ToolHandler] = {
            "shell": run_shell,
            "write_file": write_file,
            "read_file": read_file,
            "finish": finish_step,
            "flaky": run_flaky,
            "web_search": web_search,
            "web_fetch": web_fetch,
        }

    def get(self, kind: str) -> Optional[ToolHandler]:
        """Return handler for the given kind, or None if unknown."""
        return self._handlers.get(kind)