# SPDX-License-Identifier: Apache-2.0
"""
Provides the abstract base for all evaluation metrics, both built-in and user-defined custom metrics.
"""

from __future__ import annotations

import abc
from typing import TYPE_CHECKING, Any, Literal

from pydantic import BaseModel, ConfigDict, Field

if TYPE_CHECKING:
    from arksim.llms.chat.base.base_llm import BaseLLM

MetricScope = Literal["turn", "conversation"]


class ChatMessage(BaseModel):
    """A single message in a conversation."""

    role: str  # "user" or "assistant"
    content: str


def format_chat_history(messages: list[ChatMessage]) -> str:
    """Serialise a message list to the standard prompt string."""
    return "".join(f"{m.role}: {m.content}\n" for m in messages)


class QuantResult(BaseModel):
    name: str
    value: float
    reason: str | None = None
    metadata: dict[str, Any] | None = None
    scope: MetricScope = "turn"


class QualResult(BaseModel):
    """Result of a qualitative metric evaluation (e.g. agent behavior failure)."""

    name: str
    value: str
    reason: str | None = None
    metadata: dict[str, Any] | None = None
    scope: MetricScope = "turn"


class ScoreInput(BaseModel):
    """Input context provided to custom metrics for each turn.

    Extra fields are allowed — users can pass additional context
    and access it via score_input.model_extra.
    """

    model_config = ConfigDict(extra="allow")

    chat_history: list[ChatMessage] = Field(
        default_factory=list,
        description="Full conversation up to and including this turn",
    )
    current_turn: list[ChatMessage] = Field(
        default_factory=list,
        description="Only the user + assistant exchange for this turn",
    )
    knowledge: str = ""
    user_goal: str = ""
    profile: str = ""
    tool_calls: list = Field(
        default_factory=list,
        description="Tool calls observed during this turn (list of ToolCall or dict)",
    )


class _LLMMixin:
    """Mixin providing LLM injection support for metric base classes.

    Exposes a read-only ``llm`` property backed by ``self._llm``.  Subclass
    ``__init__`` is responsible for setting ``self._llm`` at construction time.
    There is no setter: injection is a load-time operation performed by the
    evaluator, not something metrics should mutate at runtime.
    """

    _llm: BaseLLM | None = None

    @property
    def llm(self) -> BaseLLM:
        if self._llm is None:
            raise RuntimeError(
                f"{self.__class__.__name__}.llm is not set. "
                "Metrics that need an LLM should accept 'llm=None' in __init__ "
                "and pass it to super().__init__(llm=llm). The evaluator injects "
                "the configured LLM automatically at load time."
            )
        return self._llm


class QuantitativeMetric(_LLMMixin, abc.ABC):
    """Abstract base class for numeric score metrics.

    Args:
        name: The name of the metric. If not provided, uses the class name as default.
        score_range: A (min, max) tuple describing the possible output range of
            :meth:`score`. Defaults to ``(1, 5)`` (the same scale as the built-in
            metrics). The range is used to normalise the value for the aggregated
            turn score and to derive a per-metric failure threshold
            (60 % of *max*).
        scope: When the metric runs: ``"turn"`` (default) evaluates every
            individual turn; ``"conversation"`` evaluates once per conversation
            alongside goal completion.
        llm: Optional LLM instance injected by the evaluator. Custom metrics that
            need an LLM for scoring should accept ``llm=None`` in their
            ``__init__`` and pass it to ``super().__init__(llm=llm)``, then use
            ``self.llm`` instead of constructing their own LLM instance.

    Example:
        >>> from arksim.evaluator import QuantitativeMetric, ScoreInput, QuantResult
        >>>
        >>> class MyMetric(QuantitativeMetric):
        >>>     def __init__(self, name: str, threshold: float = 0.5, llm=None):
        >>>         super().__init__(
        >>>             name=name,
        >>>             score_range=(0, 5),
        >>>             additional_input={"threshold": threshold},
        >>>             llm=llm,
        >>>         )
        >>>
        >>>     def score(self, score_input: ScoreInput) -> QuantResult:
        >>>         # Access standard fields: score_input.chat_history, score_input.knowledge, etc.
        >>>         # Access extra input, example: score_input.model_extra["threshold"]
        >>>
        >>>         return QuantResult(
        >>>             value=0,
        >>>             name=self.name,
        >>>             reason="Optional reason for the score"
        >>>         )
    """

    def __init__(
        self,
        name: str | None = None,
        score_range: tuple[float, float] = (1, 5),
        additional_input: dict[str, Any] | None = None,
        description: str = "",
        llm: BaseLLM | None = None,
        scope: MetricScope = "turn",
    ) -> None:
        self.name = name if name is not None else self.__class__.__name__
        self.score_range = score_range
        self.additional_input = additional_input or {}
        self.description = description
        self._llm = llm
        self.scope = scope

    def run(self, score_input: ScoreInput) -> QuantResult:
        """Score *score_input* and stamp ``scope`` on the result."""
        r = self.score(score_input.model_copy(update=self.additional_input))
        r.scope = self.scope
        return r

    @abc.abstractmethod
    def score(self, score_input: ScoreInput) -> QuantResult:
        raise NotImplementedError()


class QualitativeMetric(_LLMMixin, abc.ABC):
    """Abstract base for qualitative (categorical-label) metrics.

    Unlike QuantitativeMetric which returns a numeric score, qualitative
    metrics classify behaviour into a categorical label.

    Args:
        name: The name of the metric. If not provided, uses the class name as default.
        description: Human-readable description of what the metric measures.
        label_colors: Mapping of label values to hex colour strings for the report UI.
        scope: When the metric runs: ``"turn"`` (default) evaluates every
            individual turn; ``"conversation"`` evaluates once per conversation
            alongside goal completion.
        llm: Optional LLM instance injected by the evaluator. Custom metrics that
            need an LLM for evaluation should accept ``llm=None`` in their
            ``__init__`` and pass it to ``super().__init__(llm=llm)``, then use
            ``self.llm`` instead of constructing their own LLM instance.

    Example:
        >>> from arksim.evaluator import QualitativeMetric, QualResult, ScoreInput
        >>>
        >>> class MyQualMetric(QualitativeMetric):
        >>>     def __init__(self, llm=None):
        >>>         super().__init__(
        >>>             name="my_qual_metric",
        >>>             label_colors={
        >>>                 "pass": "#22c55e",  # green
        >>>                 "fail": "#ef4444",  # red
        >>>             },
        >>>             llm=llm,
        >>>         )
        >>>
        >>>     def evaluate(self, score_input: ScoreInput) -> QualResult:
        >>>         # Use self.llm to call the injected LLM.
        >>>         return QualResult(
        >>>             name=self.name,
        >>>             value="pass",
        >>>             reason="Optional reason",
        >>>         )
    """

    def __init__(
        self,
        name: str | None = None,
        description: str = "",
        label_colors: dict[str, str] | None = None,
        llm: BaseLLM | None = None,
        scope: MetricScope = "turn",
        additional_input: dict[str, Any] | None = None,
    ) -> None:
        self.name = name if name is not None else self.__class__.__name__
        self.description = description
        self.label_colors: dict[str, str] = label_colors or {}
        self._llm = llm
        self.scope = scope
        self.additional_input = additional_input or {}

    def run(self, score_input: ScoreInput) -> QualResult:
        """Evaluate *score_input* and stamp ``scope`` on the result."""
        r = self.evaluate(score_input.model_copy(update=self.additional_input))
        r.scope = self.scope
        return r

    @abc.abstractmethod
    def evaluate(self, score_input: ScoreInput) -> QualResult:
        raise NotImplementedError()
