from __future__ import annotations

from collections.abc import Sequence
import inspect
from typing import Callable, Generic, Protocol, TypeVar, overload, cast

from pydantic import JsonValue, TypeAdapter

from restate_agent.context import RestateRunContext
from restate_agent.errors import ToolError
from restate_agent.types import ToolHandler, ToolSpec

DepsT = TypeVar("DepsT")
OutputT = TypeVar("OutputT")
OutputT_co = TypeVar("OutputT_co", covariant=True)


class AgentRunner(Protocol[DepsT, OutputT_co]):
    async def run(
        self,
        prompt: str,
        *,
        deps: DepsT | None = None,
        tools: Sequence[ToolSpec[DepsT]] | None = None,
    ) -> OutputT_co:  # pragma: no cover - protocol
        ...


class RestateAgent(Generic[DepsT, OutputT]):
    def __init__(
        self,
        agent: AgentRunner[DepsT, OutputT],
        *,
        name: str = "restate_agent",
        output_type: type[OutputT] | None = None,
    ) -> None:
        self._agent: AgentRunner[DepsT, OutputT] = agent
        self.name: str = name
        self._tools: dict[str, ToolSpec[DepsT]] = {}
        self._output_type: type[OutputT] | None = output_type
        self._output_adapter: TypeAdapter[OutputT] | None = (
            TypeAdapter(output_type) if output_type else None
        )

    @overload
    def tool(
        self,
        func: ToolHandler[DepsT],
        *,
        name: str | None = None,
        description: str | None = None,
    ) -> ToolHandler[DepsT]:
        ...

    @overload
    def tool(
        self,
        func: None = None,
        *,
        name: str | None = None,
        description: str | None = None,
    ) -> Callable[[ToolHandler[DepsT]], ToolHandler[DepsT]]:
        ...

    def tool(
        self,
        func: ToolHandler[DepsT] | None = None,
        *,
        name: str | None = None,
        description: str | None = None,
    ) -> ToolHandler[DepsT] | Callable[[ToolHandler[DepsT]], ToolHandler[DepsT]]:
        if func is None:
            def decorator(f: ToolHandler[DepsT]) -> ToolHandler[DepsT]:
                return self.tool(f, name=name, description=description)

            return decorator

        tool_name = name or func.__name__
        if tool_name in self._tools:
            raise ToolError(f"Tool '{tool_name}' is already registered")
        tool_description = description or (inspect.getdoc(func) or "")
        self._tools[tool_name] = ToolSpec(
            name=tool_name,
            description=tool_description,
            handler=func,
        )
        return func

    @property
    def tools(self) -> Sequence[ToolSpec[DepsT]]:
        return list(self._tools.values())

    async def run(
        self,
        prompt: str,
        *,
        deps: DepsT | None = None,
        run_id: str | None = None,
        runtime: RuntimeProtocol[DepsT, OutputT] | None = None,
    ) -> OutputT:
        if runtime is None:
            runtime = cast(RuntimeProtocol[DepsT, OutputT], _default_runtime())
        return await runtime.run(self, prompt, deps, run_id)

    async def invoke_tool(
        self,
        name: str,
        payload: JsonValue,
        *,
        deps: DepsT | None = None,
        run_id: str | None = None,
        metadata: dict[str, JsonValue] | None = None,
    ) -> JsonValue:
        try:
            spec = self._tools[name]
        except KeyError as exc:
            raise ToolError(f"Tool '{name}' is not registered") from exc
        context = RestateRunContext(deps=deps, run_id=run_id, metadata=metadata or {})
        result = spec.handler(context, payload)
        if inspect.isawaitable(result):
            return await result
        return result

    async def execute(self, *, prompt: str, deps: DepsT | None, run_id: str | None) -> OutputT:
        _ = run_id
        result = await self._agent.run(prompt, deps=deps, tools=self.tools)
        return self.coerce_output(result)

    def coerce_output(self, value: object) -> OutputT:
        if self._output_adapter is None:
            return cast(OutputT, value)
        if self._output_type and isinstance(value, self._output_type):
            return value
        return self._output_adapter.validate_python(value)


class RuntimeProtocol(Protocol[DepsT, OutputT]):
    async def run(
        self,
        agent: "RestateAgent[DepsT, OutputT]",
        prompt: str,
        deps: DepsT | None,
        run_id: str | None,
    ) -> OutputT:  # pragma: no cover - protocol
        ...


class LocalRuntime(Generic[DepsT, OutputT]):
    async def run(
        self,
        agent: "RestateAgent[DepsT, OutputT]",
        prompt: str,
        deps: DepsT | None,
        run_id: str | None,
    ) -> OutputT:
        return await agent.execute(prompt=prompt, deps=deps, run_id=run_id)


def _default_runtime() -> LocalRuntime[object, object]:
    return LocalRuntime()
