"""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, TypeAlias, overload

from ..content import AssistantContentPart, Text, Thought, ToolCall
from ..exceptions import ParseError
from ..formatting import (
    Format,
    FormattableT,
    Partial,
    create_wrapper_model,
    is_output_parser,
    is_primitive_type,
)
from ..messages import Message
from ..tools import ToolkitT
from . import _utils
from .finish_reason import FinishReason
from .usage import Usage

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

AnyResponse: TypeAlias = "RootResponse[Any, Any]"


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

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

    provider_id: "ProviderId"
    """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`.
    """

    usage: Usage | None
    """Token usage statistics for this response, if available."""

    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.

        Args:
            partial: If True, parse incomplete JSON as Partial model. Only works with
                streaming responses that have accumulated JSON. Returns None if JSON
                is not yet available or cannot be parsed.

        Supports:
        - Pydantic BaseModel types (JSON schema validation)
        - Primitive types (automatically unwrapped from wrapper model)
        - Custom OutputParsers (custom parsing logic)
        - Partial parsing during streaming (when partial=True)

        Returns:
            The formatted response object of type FormatT. For BaseModel types, returns
            the model instance. For primitive types, returns the unwrapped value (e.g.,
            a string, list, dict, etc.). For OutputParsers, returns whatever the parser
            returns. When partial=True, returns None if JSON is incomplete or unparsable.

        Raises:
            NotImplementedError: If partial=True with OutputParser.
            ParseError: If parsing fails. The `original_exception` attribute contains the
                underlying error (ValueError for JSON extraction, json.JSONDecodeError
                for invalid JSON, pydantic.ValidationError for schema validation, or
                any exception from a custom OutputParser).
        """
        if self.format is None:
            return None

        formattable = self.format.formattable

        if is_output_parser(formattable):
            if partial:
                raise NotImplementedError(
                    "parse(partial=True) not supported with OutputParser. "
                    "Use BaseModel or primitive types."
                )
            try:
                return formattable(self)
            except Exception as e:
                raise ParseError(
                    f"OutputParser failed: {e}",
                    original_exception=e,
                ) from e

        if formattable is None or formattable is NoneType:  # pyright: ignore[reportUnnecessaryComparison]
            # note: pyright claims the None comparison is unnecessary, but removing it
            # introduces type errors.
            return None  # pragma: no cover

        text = self.text("")

        if partial:
            return _utils.parse_partial_json(text, formattable)
        else:
            try:
                json_text = _utils.extract_serialized_json(text)
                if is_primitive_type(formattable):
                    wrapper_model = create_wrapper_model(formattable)
                    wrapper_instance = wrapper_model.model_validate_json(json_text)
                    return wrapper_instance.output

                return formattable.model_validate_json(json_text)
            except Exception as e:
                raise ParseError(
                    f"Failed to parse response: {e}",
                    original_exception=e,
                ) from e

    def text(self, sep: str = "\n") -> str:
        """Return all text content from this response as a single string.

        Joins the text from all `Text` parts in the response content using the
        specified separator.

        Args:
            sep: The separator to use when joining multiple text parts.
                Defaults to newline ("\\n").

        Returns:
            A string containing all text content joined by the separator.
            Returns an empty string if the response contains no text parts.

        Example:
            >>> response.text()  # Join with newlines (default)
            'Hello\\nWorld'
            >>> response.text(sep=" ")  # Join with spaces
            'Hello World'
            >>> response.text(sep="")  # Concatenate directly
            'HelloWorld'
        """
        return sep.join(text.text for text in self.texts)

    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!
        """
        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 use_model  # Dynamic import to avoid circular dependency

        return use_model(self.model_id, **self.params)
