from __future__ import annotations

from collections.abc import Awaitable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable, Generic, Protocol, TypeVar

from pydantic import JsonValue

if TYPE_CHECKING:
    from restate_agent.context import RestateRunContext

DepsT = TypeVar("DepsT")
OutputT = TypeVar("OutputT")

ToolHandler = Callable[
    ["RestateRunContext[DepsT]", JsonValue],
    Awaitable[JsonValue] | JsonValue,
]


@dataclass(frozen=True)
class ToolSpec(Generic[DepsT]):
    name: str
    description: str
    handler: ToolHandler[DepsT]


class ActivityExecutor(Protocol):
    async def run(
        self,
        name: str,
        payload: dict[str, JsonValue],
    ) -> JsonValue:  # pragma: no cover - protocol
        ...
