from __future__ import annotations
from typing import Any, Awaitable, Callable, Dict, Optional
from pydantic import BaseModel, Field

ToolFn = Callable[..., Awaitable[Any]]

class ToolSpec(BaseModel):
    name: str
    description: str = ""
    input_schema: Dict[str, Any] = Field(default_factory=dict)
    output_schema: Dict[str, Any] = Field(default_factory=dict)
    tags: list[str] = Field(default_factory=list)

class Tool(BaseModel):
    spec: ToolSpec
    fn: ToolFn

class ToolRegistry:
    def __init__(self):
        self._tools: Dict[str, Tool] = {}

    def register(self, spec: ToolSpec, fn: ToolFn) -> None:
        self._tools[spec.name] = Tool(spec=spec, fn=fn)

    def get(self, name: str) -> Optional[Tool]:
        return self._tools.get(name)

    def list_specs(self) -> Dict[str, ToolSpec]:
        return {k: v.spec for k, v in self._tools.items()}
