"""Implements Response and AsyncResponse."""

import asyncio
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Generic, overload

from ..content import ToolOutput
from ..context import Context, DepsT
from ..formatting import Format, FormattableT
from ..messages import AssistantMessage, Message, UserContent
from ..tools import (
    AsyncContextTool,
    AsyncContextToolkit,
    AsyncTool,
    AsyncToolkit,
    ContextTool,
    ContextToolkit,
    Tool,
    Toolkit,
)
from .base_response import BaseResponse
from .finish_reason import FinishReason

if TYPE_CHECKING:
    from ..clients import ModelId, Params, Provider


class Response(BaseResponse[Toolkit, FormattableT]):
    """The response generated by an LLM."""

    def __init__(
        self,
        *,
        raw: Any,  # noqa: ANN401
        provider: "Provider",
        model_id: "ModelId",
        params: "Params",
        tools: Sequence[Tool] | Toolkit | None = None,
        format: Format[FormattableT] | None = None,
        input_messages: Sequence[Message],
        assistant_message: AssistantMessage,
        finish_reason: FinishReason | None,
    ) -> None:
        """Initialize a `Response`."""
        toolkit = tools if isinstance(tools, Toolkit) else Toolkit(tools=tools)
        super().__init__(
            raw=raw,
            provider=provider,
            model_id=model_id,
            params=params,
            toolkit=toolkit,
            format=format,
            input_messages=input_messages,
            assistant_message=assistant_message,
            finish_reason=finish_reason,
        )

    def execute_tools(self) -> Sequence[ToolOutput]:
        """Execute and return all of the tool calls in the response.

        Returns:
            A sequence containing a `ToolOutput` for every tool call in the order they appeared.

        Raises:
            ToolNotFoundError: If one of the response's tool calls has no matching tool.
            Exception: If one of the tools throws an exception.
        """
        return [self.toolkit.execute(tool_call) for tool_call in self.tool_calls]

    @overload
    def resume(self: "Response", content: UserContent) -> "Response": ...

    @overload
    def resume(
        self: "Response[FormattableT]", content: UserContent
    ) -> "Response[FormattableT]": ...

    def resume(self, content: UserContent) -> "Response | Response[FormattableT]":
        """Generate a new `Response` using this response's messages with additional user content.

        Uses this response's tools and format type. Also uses this response's provider,
        model, client, and params, unless the model context manager is being used to
        provide a new LLM as an override.

        Args:
            content: The new user message content to append to the message history.

        Returns:
            A new `Response` instance generated from the extended message history.
        """
        return self.model.resume(
            response=self,
            content=content,
        )


class AsyncResponse(BaseResponse[AsyncToolkit, FormattableT]):
    """The response generated by an LLM in async mode."""

    def __init__(
        self,
        *,
        raw: Any,  # noqa: ANN401
        provider: "Provider",
        model_id: "ModelId",
        params: "Params",
        tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
        format: Format[FormattableT] | None = None,
        input_messages: Sequence[Message],
        assistant_message: AssistantMessage,
        finish_reason: FinishReason | None,
    ) -> None:
        """Initialize an `AsyncResponse`."""
        toolkit = (
            tools if isinstance(tools, AsyncToolkit) else AsyncToolkit(tools=tools)
        )
        super().__init__(
            raw=raw,
            provider=provider,
            model_id=model_id,
            params=params,
            toolkit=toolkit,
            format=format,
            input_messages=input_messages,
            assistant_message=assistant_message,
            finish_reason=finish_reason,
        )

    async def execute_tools(self) -> Sequence[ToolOutput]:
        """Execute and return all of the tool calls in the response.

        Returns:
            A sequence containing a `ToolOutput` for every tool call in the order they appeared.

        Raises:
            ToolNotFoundError: If one of the response's tool calls has no matching tool.
            Exception: If one of the tools throws an exception.
        """
        tasks = [self.toolkit.execute(tool_call) for tool_call in self.tool_calls]
        return await asyncio.gather(*tasks)

    @overload
    async def resume(
        self: "AsyncResponse", content: UserContent
    ) -> "AsyncResponse": ...

    @overload
    async def resume(
        self: "AsyncResponse[FormattableT]", content: UserContent
    ) -> "AsyncResponse[FormattableT]": ...

    async def resume(
        self, content: UserContent
    ) -> "AsyncResponse | AsyncResponse[FormattableT]":
        """Generate a new `AsyncResponse` using this response's messages with additional user content.

        Uses this response's tools and format type. Also uses this response's provider,
        model, client, and params, unless the model context manager is being used to
        provide a new LLM as an override.

        Args:
            content: The new user message content to append to the message history.

        Returns:
            A new `AsyncResponse` instance generated from the extended message history.
        """
        return await self.model.resume_async(
            response=self,
            content=content,
        )


class ContextResponse(
    BaseResponse[ContextToolkit[DepsT], FormattableT], Generic[DepsT, FormattableT]
):
    """The response generated by an LLM from a context call."""

    def __init__(
        self,
        *,
        raw: Any,  # noqa: ANN401
        provider: "Provider",
        model_id: "ModelId",
        params: "Params",
        tools: Sequence[Tool | ContextTool[DepsT]]
        | ContextToolkit[DepsT]
        | None = None,
        format: Format[FormattableT] | None = None,
        input_messages: Sequence[Message],
        assistant_message: AssistantMessage,
        finish_reason: FinishReason | None,
    ) -> None:
        """Initialize a `ContextResponse`."""
        toolkit = (
            tools if isinstance(tools, ContextToolkit) else ContextToolkit(tools=tools)
        )
        super().__init__(
            raw=raw,
            provider=provider,
            model_id=model_id,
            params=params,
            toolkit=toolkit,
            format=format,
            input_messages=input_messages,
            assistant_message=assistant_message,
            finish_reason=finish_reason,
        )

    def execute_tools(self, ctx: Context[DepsT]) -> Sequence[ToolOutput]:
        """Execute and return all of the tool calls in the response.

        Args:
            ctx: A `Context` with the required deps type.

        Returns:
            A sequence containing a `ToolOutput` for every tool call.

        Raises:
            ToolNotFoundError: If one of the response's tool calls has no matching tool.
            Exception: If one of the tools throws an exception.
        """
        return [self.toolkit.execute(ctx, tool_call) for tool_call in self.tool_calls]

    @overload
    def resume(
        self: "ContextResponse[DepsT]", ctx: Context[DepsT], content: UserContent
    ) -> "ContextResponse[DepsT]": ...

    @overload
    def resume(
        self: "ContextResponse[DepsT, FormattableT]",
        ctx: Context[DepsT],
        content: UserContent,
    ) -> "ContextResponse[DepsT, FormattableT]": ...

    def resume(
        self, ctx: Context[DepsT], content: UserContent
    ) -> "ContextResponse[DepsT] | ContextResponse[DepsT, FormattableT]":
        """Generate a new `ContextResponse` using this response's messages with additional user content.

        Uses this response's tools and format type. Also uses this response's provider,
        model, client, and params, unless the model context manager is being used to
        provide a new LLM as an override.

        Args:
            ctx: A `Context` with the required deps type.
            content: The new user message content to append to the message history.

        Returns:
            A new `ContextResponse` instance generated from the extended message history.
        """
        return self.model.context_resume(
            ctx=ctx,
            response=self,
            content=content,
        )


class AsyncContextResponse(
    BaseResponse[AsyncContextToolkit[DepsT], FormattableT], Generic[DepsT, FormattableT]
):
    """The response generated by an LLM from an async context call."""

    def __init__(
        self,
        *,
        raw: Any,  # noqa: ANN401
        provider: "Provider",
        model_id: "ModelId",
        params: "Params",
        tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
        | AsyncContextToolkit[DepsT]
        | None = None,
        format: Format[FormattableT] | None = None,
        input_messages: Sequence[Message],
        assistant_message: AssistantMessage,
        finish_reason: FinishReason | None,
    ) -> None:
        """Initialize an `AsyncContextResponse`."""
        toolkit = (
            tools
            if isinstance(tools, AsyncContextToolkit)
            else AsyncContextToolkit(tools=tools)
        )
        super().__init__(
            raw=raw,
            provider=provider,
            model_id=model_id,
            params=params,
            toolkit=toolkit,
            format=format,
            input_messages=input_messages,
            assistant_message=assistant_message,
            finish_reason=finish_reason,
        )

    async def execute_tools(self, ctx: Context[DepsT]) -> Sequence[ToolOutput]:
        """Execute and return all of the tool calls in the response.

        Args:
            ctx: A `Context` with the required deps type.

        Returns:
            A sequence containing a `ToolOutput` for every tool call in the order they appeared.

        Raises:
            ToolNotFoundError: If one of the response's tool calls has no matching tool.
            Exception: If one of the tools throws an exception.
        """
        tasks = [self.toolkit.execute(ctx, tool_call) for tool_call in self.tool_calls]
        return await asyncio.gather(*tasks)

    @overload
    async def resume(
        self: "AsyncContextResponse[DepsT]", ctx: Context[DepsT], content: UserContent
    ) -> "AsyncContextResponse[DepsT]": ...

    @overload
    async def resume(
        self: "AsyncContextResponse[DepsT, FormattableT]",
        ctx: Context[DepsT],
        content: UserContent,
    ) -> "AsyncContextResponse[DepsT, FormattableT]": ...

    async def resume(
        self, ctx: Context[DepsT], content: UserContent
    ) -> "AsyncContextResponse[DepsT] | AsyncContextResponse[DepsT, FormattableT]":
        """Generate a new `AsyncContextResponse` using this response's messages with additional user content.

        Uses this response's tools and format type. Also uses this response's provider,
        model, client, and params, unless the model context manager is being used to
        provide a new LLM as an override.

        Args:
            ctx: A Context with the required deps type.
            content: The new user message content to append to the message history.

        Returns:
            A new `AsyncContextResponse` instance generated from the extended message history.
        """
        return await self.model.context_resume_async(
            ctx=ctx,
            response=self,
            content=content,
        )
