"""
Tool registry for workflow engine steps.
"""
from __future__ import annotations

from typing import Any, Awaitable, Callable, Dict, Optional

try:
    from backend.models.agent import TurnContext, ToolSpec
except ImportError:
    from models.agent import TurnContext, ToolSpec

ToolHandler = Callable[[TurnContext], Awaitable[None]]


class ToolRegistry:
    def __init__(self) -> None:
        self._tools: Dict[str, ToolHandler] = {}
        self._specs: Dict[str, ToolSpec] = {}

    def register(
        self,
        name: str,
        handler: ToolHandler,
        *,
        description: str = "",
        parameters: Optional[Dict[str, Any]] = None,
    ) -> None:
        self._tools[name] = handler
        self._specs[name] = ToolSpec(
            name=name,
            description=description,
            parameters=parameters or {},
        )

    def get(self, name: str) -> ToolHandler:
        if name not in self._tools:
            raise KeyError(f"Unknown tool: {name}")
        return self._tools[name]

    def list_specs(self) -> list[ToolSpec]:
        return list(self._specs.values())

    async def execute(self, name: str, ctx: TurnContext) -> None:
        handler = self.get(name)
        await handler(ctx)


default_tool_registry = ToolRegistry()