"""Base interface for all LLM responses."""

from abc import ABC
from collections.abc import Sequence
from types import NoneType
from typing import TYPE_CHECKING, Any, Generic, Literal, overload

from ..content import AssistantContentPart, Text, Thought, ToolCall
from ..formatting import Format, FormattableT, Partial
from ..messages import Message
from ..tools import ToolkitT
from . import _utils
from .finish_reason import FinishReason

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


class RootResponse(Generic[ToolkitT, FormattableT], ABC):
    """Base class for LLM responses."""

    raw: Any
    """The raw response from the LLM."""

    provider: "Provider"
    """The provider that generated this response."""

    model_id: "ModelId"
    """The model id that generated this response."""

    params: "Params"
    """The params that were used to generate this response (or None)."""

    toolkit: ToolkitT
    """The toolkit containing the tools used when generating this response."""

    messages: list[Message]
    """The message history, including the most recent assistant message."""

    content: Sequence[AssistantContentPart]
    """The content generated by the LLM."""

    texts: Sequence[Text]
    """The text content in the generated response, if any."""

    tool_calls: Sequence[ToolCall]
    """The tools the LLM wants called on its behalf, if any."""

    thoughts: Sequence[Thought]
    """The readable thoughts from the model's thinking process, if any.

    The thoughts may be direct output from the model thinking process, or may be a
    generated summary. (This depends on the provider; newer models tend to summarize.)
    """
    finish_reason: FinishReason | None
    """The reason why the LLM finished generating a response, if set.
    
    `finish_reason` is only set if the response did not finish generating normally,
    e.g. `FinishReason.MAX_TOKENS` if the model ran out of tokens before completing.
    When the response generates normally, `response.finish_reason` will be `None`.
    """

    format: Format[FormattableT] | None
    """The `Format` describing the structured response format, if available."""

    @overload
    def parse(self: "RootResponse[ToolkitT, None]", partial: Literal[True]) -> None:
        """Format the response into a `Partial[BaseModel]` (with optional fields).

        This is useful for when the stream is only partially consumed, in which case the
        structured output may only be partially available.
        """
        ...

    @overload
    def parse(
        self: "RootResponse[ToolkitT, FormattableT]", partial: Literal[True]
    ) -> Partial[FormattableT]:
        """Format the response into a `Partial[BaseModel]` (with optional fields).

        This is useful for when the stream is only partially consumed, in which case the
        structured output may only be partially available.
        """
        ...

    @overload
    def parse(
        self: "RootResponse[ToolkitT, None]", partial: Literal[False] = False
    ) -> None:
        """Overload when the format type is `None`."""
        ...

    @overload
    def parse(
        self: "RootResponse[ToolkitT, FormattableT]", partial: Literal[False] = False
    ) -> FormattableT:
        """Overload when the format type is not `None`."""
        ...

    def parse(
        self, partial: bool = False
    ) -> FormattableT | Partial[FormattableT] | None:
        """Format the response according to the response format parser.

        Returns:
            The formatted response object of type FormatT.

        Raises:
            json.JSONDecodeError: If the response's textual content can't be parsed as
                JSON.
            pydantic.ValidationError: If the response's content fails validation for the
                format type.
        """
        if self.format is None:
            return None

        formattable = self.format.formattable
        if formattable is None or formattable is NoneType:
            return None  # pragma: no cover

        if partial:
            raise NotImplementedError

        text = "".join(text.text for text in self.texts)
        json_text = _utils.extract_serialized_json(text)

        return formattable.model_validate_json(json_text)

    def pretty(self) -> str:
        """Return a string representation of all response content.

        The response content will be represented in a way that emphasies clarity and
        readability, but may not include all metadata (like thinking signatures or tool
        call ids), and thus cannot be used to reconstruct the response. For example:

        **Thinking:**
          The user is asking a math problem. I should use the calculator tool.

        **Tool Call (calculator)** {'operation': 'mult', 'a': 1337, 'b': 4242}

        I am going to use the calculator and answer your question for you!
        """
        if not self.content:
            return "**[No Content]**"

        pretty_parts: list[str] = []
        for part in self.content:
            if isinstance(part, Text):
                pretty_parts.append(part.text)
            elif isinstance(part, ToolCall):
                pretty_parts.append(f"**ToolCall ({part.name}):** {part.args}")
            elif isinstance(part, Thought):
                indented_thinking = "\n".join(
                    f"  {line}" for line in part.thought.split("\n")
                )
                pretty_parts.append(f"**Thinking:**\n{indented_thinking}")
            else:
                pretty_parts.append(
                    f"[{type(part).__name__}: {str(part)}]"
                )  # pragma: no cover

        return "\n\n".join(pretty_parts)

    @property
    def model(self) -> "Model":
        """A `Model` with parameters matching this response."""
        from ..models import Model, get_model_from_context

        if context_model := get_model_from_context():
            return context_model

        return Model(
            provider=self.provider,
            model_id=self.model_id,
            **self.params,
        )
