import asyncio
from collections.abc import Awaitable, Callable

from typing_extensions import override

from pydantic import BaseModel, JsonValue

from restate_agent import RestateAgent
from restate_agent.runtime import create_durable_runtime
from restate_agent.types import ActivityExecutor


class Answer(BaseModel):
    text: str


class DummyAgent:
    async def run(
        self,
        prompt: str,
        deps: None | object = None,
        tools: object | None = None,
    ) -> Answer:
        _ = deps
        _ = tools
        return Answer(text=f"Durable echo: {prompt}")


class LocalExecutor(ActivityExecutor):
    def __init__(self, handler: Callable[[dict[str, JsonValue]], Awaitable[JsonValue]]) -> None:
        self._handler: Callable[[dict[str, JsonValue]], Awaitable[JsonValue]] = handler

    @override
    async def run(self, name: str, payload: dict[str, JsonValue]) -> JsonValue:
        _ = name
        return await self._handler(payload)


async def main() -> None:
    agent: RestateAgent[None, Answer] = RestateAgent(DummyAgent(), output_type=Answer)
    runtime, handler = create_durable_runtime(agent, LocalExecutor)
    _ = handler
    result = await agent.run("hello", runtime=runtime)
    print(result.text)


if __name__ == "__main__":
    asyncio.run(main())
