"""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 ..exceptions import ParseError
from ..formatting import Format, FormattableT
from ..messages import AssistantMessage, Message, UserContent
from ..tools import (
    AsyncContextToolkit,
    AsyncContextTools,
    AsyncToolkit,
    AsyncTools,
    ContextToolkit,
    ContextTools,
    Toolkit,
    Tools,
)
from ..types import Jsonable
from .base_response import BaseResponse
from .finish_reason import FinishReason
from .usage import Usage

if TYPE_CHECKING:
    from ..models import Params
    from ..providers import ModelId, ProviderId


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

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

    def execute_tools(self) -> Sequence[ToolOutput[Jsonable]]:
        """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,
        )

    @overload
    def validate(
        self: "Response[None]", max_retries: int = 1
    ) -> tuple[None, "Response[None]"]: ...

    @overload
    def validate(
        self: "Response[FormattableT]", max_retries: int = 1
    ) -> tuple[FormattableT, "Response[FormattableT]"]: ...

    def validate(
        self, max_retries: int = 1
    ) -> tuple[FormattableT | None, "Response[FormattableT] | Response[None]"]:
        """Parse and validate the response, retrying on parse errors.

        Attempts to parse the response. On `ParseError`, asks the LLM to fix its
        output by resuming with the error message. Returns both the parsed value
        and the (potentially updated) response.

        Args:
            max_retries: Maximum number of retry attempts on parse failure.
                Defaults to 1 (2 total attempts). Must be non-negative.

        Returns:
            A tuple of (parsed_value, response). If parsing succeeded on the first
            attempt, returns (value, self). If retries were needed, returns
            (value, new_response) where new_response is the final successful response.

        Raises:
            ValueError: If max_retries is negative.
            ParseError: If parsing fails after exhausting all retry attempts.
            Error: If the LLM call fails while generating a retry response.
        """
        if max_retries < 0:
            raise ValueError("max_retries must be non-negative")

        if self.format is None:
            return None, self

        current_response: Response[FormattableT] = self
        for attempt in range(max_retries + 1):
            try:
                return current_response.parse(), current_response
            except ParseError as e:
                if attempt == max_retries:
                    raise
                current_response = current_response.resume(e.retry_message())

        raise AssertionError("Unreachable")  # pragma: no cover


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

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

    async def execute_tools(self) -> Sequence[ToolOutput[Jsonable]]:
        """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,
        )

    @overload
    async def validate(
        self: "AsyncResponse[None]", max_retries: int = 1
    ) -> tuple[None, "AsyncResponse[None]"]: ...

    @overload
    async def validate(
        self: "AsyncResponse[FormattableT]", max_retries: int = 1
    ) -> tuple[FormattableT, "AsyncResponse[FormattableT]"]: ...

    async def validate(
        self, max_retries: int = 1
    ) -> tuple[
        FormattableT | None, "AsyncResponse[FormattableT] | AsyncResponse[None]"
    ]:
        """Parse and validate the response, retrying on parse errors.

        Attempts to parse the response. On `ParseError`, asks the LLM to fix its
        output by resuming with the error message. Returns both the parsed value
        and the (potentially updated) response.

        Args:
            max_retries: Maximum number of retry attempts on parse failure.
                Defaults to 1 (2 total attempts). Must be non-negative.

        Returns:
            A tuple of (parsed_value, response). If parsing succeeded on the first
            attempt, returns (value, self). If retries were needed, returns
            (value, new_response) where new_response is the final successful response.

        Raises:
            ValueError: If max_retries is negative.
            ParseError: If parsing fails after exhausting all retry attempts.
            Error: If the LLM call fails while generating a retry response.
        """
        if max_retries < 0:
            raise ValueError("max_retries must be non-negative")

        if self.format is None:
            return None, self

        current_response: AsyncResponse[FormattableT] = self
        for attempt in range(max_retries + 1):
            try:
                return current_response.parse(), current_response
            except ParseError as e:
                if attempt == max_retries:
                    raise
                current_response = await current_response.resume(e.retry_message())

        raise AssertionError("Unreachable")  # pragma: no cover


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_id: "ProviderId",
        model_id: "ModelId",
        provider_model_name: str,
        params: "Params",
        tools: ContextTools[DepsT] | None = None,
        format: Format[FormattableT] | None = None,
        input_messages: Sequence[Message],
        assistant_message: AssistantMessage,
        finish_reason: FinishReason | None,
        usage: Usage | None,
    ) -> None:
        """Initialize a `ContextResponse`."""
        toolkit = (
            tools if isinstance(tools, ContextToolkit) else ContextToolkit(tools=tools)
        )
        super().__init__(
            raw=raw,
            provider_id=provider_id,
            model_id=model_id,
            provider_model_name=provider_model_name,
            params=params,
            toolkit=toolkit,
            format=format,
            input_messages=input_messages,
            assistant_message=assistant_message,
            finish_reason=finish_reason,
            usage=usage,
        )

    def execute_tools(self, ctx: Context[DepsT]) -> Sequence[ToolOutput[Jsonable]]:
        """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,
        )

    @overload
    def validate(
        self: "ContextResponse[DepsT, None]", ctx: Context[DepsT], max_retries: int = 1
    ) -> tuple[None, "ContextResponse[DepsT, None]"]: ...

    @overload
    def validate(
        self: "ContextResponse[DepsT, FormattableT]",
        ctx: Context[DepsT],
        max_retries: int = 1,
    ) -> tuple[FormattableT, "ContextResponse[DepsT, FormattableT]"]: ...

    def validate(
        self, ctx: Context[DepsT], max_retries: int = 1
    ) -> tuple[
        FormattableT | None,
        "ContextResponse[DepsT, FormattableT] | ContextResponse[DepsT, None]",
    ]:
        """Parse and validate the response, retrying on parse errors.

        Attempts to parse the response. On `ParseError`, asks the LLM to fix its
        output by resuming with the error message. Returns both the parsed value
        and the (potentially updated) response.

        Args:
            ctx: A `Context` with the required deps type.
            max_retries: Maximum number of retry attempts on parse failure.
                Defaults to 1 (2 total attempts). Must be non-negative.

        Returns:
            A tuple of (parsed_value, response). If parsing succeeded on the first
            attempt, returns (value, self). If retries were needed, returns
            (value, new_response) where new_response is the final successful response.

        Raises:
            ValueError: If max_retries is negative.
            ParseError: If parsing fails after exhausting all retry attempts.
            Error: If the LLM call fails while generating a retry response.
        """
        if max_retries < 0:
            raise ValueError("max_retries must be non-negative")

        if self.format is None:
            return None, self

        current_response: ContextResponse[DepsT, FormattableT] = self
        for attempt in range(max_retries + 1):
            try:
                return current_response.parse(), current_response
            except ParseError as e:
                if attempt == max_retries:
                    raise
                current_response = current_response.resume(ctx, e.retry_message())

        raise AssertionError("Unreachable")  # pragma: no cover


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_id: "ProviderId",
        model_id: "ModelId",
        provider_model_name: str,
        params: "Params",
        tools: AsyncContextTools[DepsT] | None = None,
        format: Format[FormattableT] | None = None,
        input_messages: Sequence[Message],
        assistant_message: AssistantMessage,
        finish_reason: FinishReason | None,
        usage: Usage | None,
    ) -> None:
        """Initialize an `AsyncContextResponse`."""
        toolkit = (
            tools
            if isinstance(tools, AsyncContextToolkit)
            else AsyncContextToolkit(tools=tools)
        )
        super().__init__(
            raw=raw,
            provider_id=provider_id,
            model_id=model_id,
            provider_model_name=provider_model_name,
            params=params,
            toolkit=toolkit,
            format=format,
            input_messages=input_messages,
            assistant_message=assistant_message,
            finish_reason=finish_reason,
            usage=usage,
        )

    async def execute_tools(
        self, ctx: Context[DepsT]
    ) -> Sequence[ToolOutput[Jsonable]]:
        """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,
        )

    @overload
    async def validate(
        self: "AsyncContextResponse[DepsT, None]",
        ctx: Context[DepsT],
        max_retries: int = 1,
    ) -> tuple[None, "AsyncContextResponse[DepsT, None]"]: ...

    @overload
    async def validate(
        self: "AsyncContextResponse[DepsT, FormattableT]",
        ctx: Context[DepsT],
        max_retries: int = 1,
    ) -> tuple[FormattableT, "AsyncContextResponse[DepsT, FormattableT]"]: ...

    async def validate(
        self, ctx: Context[DepsT], max_retries: int = 1
    ) -> tuple[
        FormattableT | None,
        "AsyncContextResponse[DepsT, FormattableT] | AsyncContextResponse[DepsT, None]",
    ]:
        """Parse and validate the response, retrying on parse errors.

        Attempts to parse the response. On `ParseError`, asks the LLM to fix its
        output by resuming with the error message. Returns both the parsed value
        and the (potentially updated) response.

        Args:
            ctx: A `Context` with the required deps type.
            max_retries: Maximum number of retry attempts on parse failure.
                Defaults to 1 (2 total attempts). Must be non-negative.

        Returns:
            A tuple of (parsed_value, response). If parsing succeeded on the first
            attempt, returns (value, self). If retries were needed, returns
            (value, new_response) where new_response is the final successful response.

        Raises:
            ValueError: If max_retries is negative.
            ParseError: If parsing fails after exhausting all retry attempts.
            Error: If the LLM call fails while generating a retry response.
        """
        if max_retries < 0:
            raise ValueError("max_retries must be non-negative")

        if self.format is None:
            return None, self

        current_response: AsyncContextResponse[DepsT, FormattableT] = self
        for attempt in range(max_retries + 1):
            try:
                return current_response.parse(), current_response
            except ParseError as e:
                if attempt == max_retries:
                    raise
                current_response = await current_response.resume(ctx, e.retry_message())

        raise AssertionError("Unreachable")  # pragma: no cover
