"""Base abstractions for the AgentByte evaluation framework."""

from __future__ import annotations

from abc import ABC, abstractmethod

from agentbyte.cancellation_token import CancellationToken
from agentbyte.messages import AssistantMessage

from .types import AnswerStrategy, EvalScore, EvalTask, EvalTrajectory

_VALID_ANSWER_STRATEGIES: tuple[AnswerStrategy, ...] = (
    "last_non_empty",
    "last_assistant",
    "last_content",
    "all_assistant",
)


class BaseEvalTarget(ABC):
    """Anything that can execute an evaluation task into a trajectory."""

    def __init__(self, name: str):
        self.name = name

    @abstractmethod
    async def run(
        self,
        task: EvalTask,
        cancellation_token: CancellationToken | None = None,
    ) -> EvalTrajectory:
        """Execute an evaluation task and return its trajectory."""


class BaseEvalJudge(ABC):
    """Base class for evaluation judges."""

    def __init__(
        self,
        name: str,
        answer_strategy: AnswerStrategy = "last_non_empty",
        source_filter: str | None = None,
    ) -> None:
        if answer_strategy not in _VALID_ANSWER_STRATEGIES:
            raise ValueError(
                f"Unsupported answer_strategy: {answer_strategy!r}. "
                f"Expected one of {list(_VALID_ANSWER_STRATEGIES)}."
            )
        self.name = name
        self.answer_strategy = answer_strategy
        self.source_filter = source_filter

    def extract_answer(self, trajectory: EvalTrajectory) -> str:
        """Extract the answer text from a trajectory using the configured strategy.

        When ``source_filter`` is set on this judge, messages are pre-filtered to
        only those whose ``source`` matches before the strategy is applied. This
        lets callers target a specific agent's output in multi-agent trajectories
        (e.g. extract the Writer's last message from a Researcher→Writer→Reviewer
        orchestrator run).
        """
        messages = list(trajectory.messages)
        if self.source_filter:
            messages = [m for m in messages if m.source == self.source_filter]
        if not messages:
            return ""

        if self.answer_strategy == "last_content":
            return messages[-1].content

        if self.answer_strategy == "last_non_empty":
            for message in reversed(messages):
                if message.content and message.content.strip():
                    return message.content
            return ""

        if self.answer_strategy == "last_assistant":
            for message in reversed(messages):
                if isinstance(message, AssistantMessage):
                    return message.content
            return ""

        assistant_contents = [
            message.content
            for message in messages
            if isinstance(message, AssistantMessage)
            and message.content
            and message.content.strip()
        ]
        return "\n".join(assistant_contents)

    @abstractmethod
    async def score(
        self,
        trajectory: EvalTrajectory,
        criteria: list[str] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> EvalScore:
        """Score an evaluation trajectory."""


class BaseEvalRunner(ABC):
    """Base class for evaluation runners."""

    def __init__(self, judge: BaseEvalJudge):
        self.judge = judge

    @abstractmethod
    async def evaluate(
        self,
        target: BaseEvalTarget,
        tasks: list[EvalTask],
        criteria: list[str] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> list[EvalScore]:
        """Evaluate a target on multiple tasks."""


__all__ = ["BaseEvalTarget", "BaseEvalJudge", "BaseEvalRunner"]
