import asyncio
from collections.abc import Awaitable, Callable
from typing import cast

from typing_extensions import override

from pydantic import BaseModel, JsonValue
import pytest

from restate_agent import RestateAgent, RestateRunContext
from restate_agent.errors import RestateRuntimeError, ToolError
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"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)


def test_agent_run_coerces_output():
    async def run():
        agent: RestateAgent[None, Answer] = RestateAgent(DummyAgent(), output_type=Answer)
        result = await agent.run("hello")
        assert isinstance(result, Answer)
        assert result.text == "Echo: hello"

    asyncio.run(run())


def test_agent_tools_invoke():
    async def run():
        agent: RestateAgent[None, Answer] = RestateAgent(DummyAgent(), output_type=Answer)

        @agent.tool
        async def add(ctx: RestateRunContext[None], payload: JsonValue) -> int:
            _ = ctx
            payload_dict = cast(dict[str, int], payload)
            return payload_dict["a"] + payload_dict["b"]
        _ = add

        result = await agent.invoke_tool("add", {"a": 2, "b": 3}, deps=None)
        assert result == 5

    asyncio.run(run())


def test_tool_registration_errors():
    agent: RestateAgent[None, Answer] = RestateAgent(DummyAgent(), output_type=Answer)

    @agent.tool(name="dupe")
    async def first(ctx: RestateRunContext[None], payload: JsonValue) -> int:
        _ = ctx
        _ = payload
        return 1
    _ = first

    with pytest.raises(ToolError):
        _ = agent.tool(first, name="dupe")

    async def run():
        with pytest.raises(ToolError):
            _ = await agent.invoke_tool("missing", {"a": 1}, deps=None)

    asyncio.run(run())


def test_durable_runtime_executes_activity():
    async def run():
        agent: RestateAgent[None, Answer] = RestateAgent(DummyAgent(), output_type=Answer)
        runtime, handler = create_durable_runtime(agent, LocalExecutor)
        _ = handler
        result = await agent.run("durable", runtime=runtime)
        assert result.text == "Echo: durable"

    asyncio.run(run())


def test_durable_runtime_rejects_non_json_deps():
    async def run():
        agent: RestateAgent[object, Answer] = RestateAgent(DummyAgent(), output_type=Answer)
        runtime, _handler = create_durable_runtime(agent, LocalExecutor)
        _ = _handler

        class NonJsonDeps:
            pass

        with pytest.raises(RestateRuntimeError):
            _ = await agent.run("bad-deps", deps=NonJsonDeps(), runtime=runtime)

    asyncio.run(run())


def test_durable_runtime_rejects_non_json_output():
    class NonJson:
        pass

    class NonJsonAgent:
        async def run(
            self,
            prompt: str,
            deps: JsonValue | None = None,
            tools: object | None = None,
        ) -> NonJson:
            _ = prompt
            _ = deps
            _ = tools
            return NonJson()

    async def run():
        agent: RestateAgent[None, NonJson] = RestateAgent(NonJsonAgent())
        runtime, _handler = create_durable_runtime(agent, LocalExecutor)
        _ = _handler
        with pytest.raises(RestateRuntimeError):
            _ = await agent.run("bad", runtime=runtime)

    asyncio.run(run())
