from __future__ import annotations

from collections.abc import Awaitable
from typing import TYPE_CHECKING, Callable, Generic, TypeVar, cast

from typing_extensions import override

from pydantic import BaseModel, JsonValue, TypeAdapter

from restate_agent.errors import RestateRuntimeError
from restate_agent.types import ActivityExecutor

if TYPE_CHECKING:
    from restate_agent.agent import RestateAgent

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

_JSON_ADAPTER: TypeAdapter[JsonValue] = TypeAdapter(JsonValue)


class AgentRunPayload(BaseModel):
    prompt: str
    deps: JsonValue | None = None
    run_id: str | None = None


class RestateActivityExecutor(ActivityExecutor):
    def __init__(self, context: object, *, method_name: str = "run_activity") -> None:
        self._context: object = context
        self._method_name: str = method_name

    @override
    async def run(self, name: str, payload: dict[str, JsonValue]) -> JsonValue:
        method = getattr(self._context, self._method_name, None)
        if method is None:
            raise RestateRuntimeError(
                f"Restate context missing '{self._method_name}' for activity execution"
            )
        runner = cast(Callable[[str, dict[str, JsonValue]], Awaitable[JsonValue]], method)
        return await runner(name, payload)


class DurableRuntime(Generic[DepsT, OutputT]):
    def __init__(self, executor: ActivityExecutor, *, activity_name: str = "restate_agent.run") -> None:
        self._executor: ActivityExecutor = executor
        self._activity_name: str = activity_name

    async def run(
        self,
        agent: "RestateAgent[DepsT, OutputT]",
        prompt: str,
        deps: DepsT | None,
        run_id: str | None,
    ) -> OutputT:
        if deps is not None and not isinstance(deps, (dict, list, str, int, float, bool)):
            raise RestateRuntimeError("durable runtime requires JSON-serializable deps")
        if deps is not None:
            _ = _JSON_ADAPTER.validate_python(deps)
        payload = AgentRunPayload(
            prompt=prompt,
            deps=cast(JsonValue | None, deps),
            run_id=run_id,
        )
        result = await self._executor.run(self._activity_name, payload.model_dump())
        return agent.coerce_output(result)


def create_durable_runtime(
    agent: "RestateAgent[DepsT, OutputT]",
    executor_factory: Callable[
        [Callable[[dict[str, JsonValue]], Awaitable[JsonValue]]],
        ActivityExecutor,
    ],
    *,
    activity_name: str = "restate_agent.run",
    ) -> tuple[DurableRuntime[DepsT, OutputT], Callable[[dict[str, JsonValue]], Awaitable[JsonValue]]]:

    async def handler(payload: dict[str, JsonValue]) -> JsonValue:
        parsed = AgentRunPayload.model_validate(payload)
        result = await agent.execute(
            prompt=parsed.prompt,
            deps=cast(DepsT | None, parsed.deps),
            run_id=parsed.run_id,
        )
        if isinstance(result, BaseModel):
            output = result.model_dump()
        else:
            output = result
        if not isinstance(output, (dict, list, str, int, float, bool, type(None))):
            raise RestateRuntimeError("durable runtime requires JSON-serializable output")
        _ = _JSON_ADAPTER.validate_python(output)
        return cast(JsonValue, output)

    runtime: DurableRuntime[DepsT, OutputT] = DurableRuntime(
        executor_factory(handler),
        activity_name=activity_name,
    )
    return runtime, handler
