"""zyx.result"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import (
    Any,
    Generic,
    List,
    TypeVar,
)

from ._aliases import (
    PydanticAIAgentResult,
    PydanticAIUsage,
    PydanticAIMessage,
)
from ._utils._confidence import Confidence, score_confidence
from ._utils._semantic import semantic_from_output, semantic_for_operation

__all__ = ("Result",)


Output = TypeVar("Output")


@dataclass
class Result(Generic[Output]):
    """
    Universal representation of the output/result of a semantic operation.

    This class is used to wrap the final `output`, which provides either a
    complete or partial representation of the `target` type or value, based on
    the operation used.
    """

    output: Output
    """
    The final output content generated by the semantic operation.

    This is either a complete, partial or mutated representation of the `target`
    type or value."""

    raw: List[PydanticAIAgentResult[Any]] = field(repr=False)
    """All underlying AgentRunResult(s) that were used to produce this result."""

    _semantic_message: str | None = field(default=None, init=False)

    @property
    def confidence(self) -> Confidence:
        """
        Confidence scores for the `output` of this result.

        NOTE: This may return `None` if `confidence` was not enabled when running the
        semantic operation this result was produced from, or if the model API does not support
        retrieving log-probabilities.

        Returns:
            A `Confidence` object containing the confidence scores for the `output` of this result.
        """
        return score_confidence(
            runs=self.raw,
            output=self.output,
            model_name=self.model,
        )  # type: ignore

    @property
    def model(self) -> str:
        """
        The model that was used to produce this result.
        """
        return self.raw[0].response.model_name or ""

    @property
    def usage(self) -> PydanticAIUsage:
        """
        The total accumulated token usage across all agent runs that were used internally
        to produce this result.
        """
        total = PydanticAIUsage()
        for run in self.raw:
            total.incr(run.usage())
        return total

    @property
    def all_new_messages(self) -> List[PydanticAIMessage]:
        """
        All new messages genenerated by `pydantic_ai` that were added/utilized during the
        execution of the operation that was used to produce this result.
        """
        messages = []

        for run in self.raw:
            messages.extend(run.new_messages())

        return messages

    def __repr__(self) -> str:
        return (
            f"Result({type(self.output).__name__}):\n"
            f"{self.output}\n\n"
            f">>> Model: {self.model}\n"
        )

    def as_message(
        self,
        *,
        operation: str | None = None,
        original: Any | None = None,
        updated: Any | None = None,
        summary: str | None = None,
    ) -> str:
        """Return a semantic message for this result."""
        if self._semantic_message:
            return self._semantic_message
        if operation:
            return semantic_for_operation(
                operation,
                output=self.output,
                original=original,
                updated=updated,
                summary=summary,
            )
        return semantic_from_output(self.output)

    def __rich__(self):
        from rich.console import RenderableType, Group
        from rich.rule import Rule
        from rich.text import Text

        renderables: list[RenderableType] = []

        renderables.append(
            Rule(
                title=f"✨ Result({type(self.output).__name__})",
                style="rule.line",
                align="left",
            )
        )

        output_text = f"{self.output}\n"
        renderables.append(Text.from_markup(f"[bold]{output_text}[/bold]"))

        renderables.append(
            Text.from_markup(
                f"[sandy_brown]>>>[/sandy_brown] [dim italic]Type: {type(self.output).__name__}[/dim italic]"
            )
        )

        model = getattr(self, "model", None)
        if model:
            model_names = (
                model
                if isinstance(model, str)
                else (model[0] if len(model) == 1 else ", ".join(model))
            )
            label = (
                "Model"
                if (
                    isinstance(model_names, str)
                    or (hasattr(model, "__len__") and len(model) == 1)
                )
                else "Models"
            )
            model_info = Text.from_markup(
                f"[sandy_brown]>>>[/sandy_brown] [dim italic]{label}: {model_names}[/dim italic]"
            )
            renderables.append(model_info)

        return Group(*renderables)
