"""Base class for self-ask toolset evaluations."""

from typing import Any

from pydantic_ai_toolsets.evals.base import EvaluationResult, ToolsetEvaluator
from pydantic_ai_toolsets.evals.config import EvaluationConfig
from pydantic_ai_toolsets.evals.datasets.self_ask_cases import SELF_ASK_CASES
from pydantic_ai_toolsets.evals.evaluators import (
    EfficiencyEvaluator,
    OutputQualityEvaluator,
    StorageStateEvaluator,
    ToolUsageEvaluator,
)


class BaseSelfAskEvaluator(ToolsetEvaluator[str, Any]):
    """Base evaluator for self-ask toolsets."""

    def __init__(
        self,
        toolset_name: str,
        config: EvaluationConfig | None = None,
    ):
        """Initialize base self-ask evaluator."""
        super().__init__(toolset_name, config)

    async def evaluate_all_cases(self) -> list[EvaluationResult]:
        """Evaluate all self-ask test cases."""
        results = []
        output_evaluator = OutputQualityEvaluator()
        tool_evaluator = ToolUsageEvaluator()
        storage_evaluator = StorageStateEvaluator()
        efficiency_evaluator = EfficiencyEvaluator()

        for case in SELF_ASK_CASES:
            result = await self.evaluate_case(case.prompt)

            # Add scores from evaluators
            quality_scores = output_evaluator.evaluate_all(
                result, case.prompt, required_aspects=None
            )
            tool_scores = tool_evaluator.evaluate_all(
                result,
                expected_sequence=None,
                required_tools=case.expected_tools,
            )
            storage_scores = storage_evaluator.evaluate_all(
                result,
                expected_keys=case.expected_storage_keys,
                min_items=case.min_storage_items,
            )
            efficiency_metrics = efficiency_evaluator.evaluate_all(result)

            # Validate depth limit
            depth_validation_score = self._validate_depth_limit(
                result, case.max_depth_expected
            )

            # Combine scores
            result.scores.update(quality_scores)
            result.scores.update(tool_scores)
            result.scores.update(storage_scores)
            result.scores["efficiency"] = efficiency_metrics.get("efficiency_score", 0.0)
            result.scores["depth_validation"] = depth_validation_score

            result.case_name = case.name
            results.append(result)

        return results

    def _validate_depth_limit(
        self, result: EvaluationResult, max_depth_expected: int
    ) -> float:
        """Validate that depth limit is respected.

        Args:
            result: Evaluation result with storage state.
            max_depth_expected: Expected maximum depth from test case.

        Returns:
            Score (0.0-1.0) based on depth validation.
        """
        if not result.storage_state:
            return 0.0

        # Extract max depth from storage statistics if available
        stats = result.storage_state.get("statistics", {})
        max_depth_reached = stats.get("max_depth_reached", 0)

        # Check if depth exceeds maximum (3)
        if max_depth_reached > 3:
            return 0.0  # Failed - exceeded maximum depth

        # Check if depth matches expected (within reason)
        # Allow some flexibility - if expected is 2, actual can be 1-3
        if max_depth_expected == 0:
            # Simple case - should stay at depth 0 or 1
            if max_depth_reached <= 1:
                return 1.0
            else:
                return 0.5  # Partial credit
        elif max_depth_expected <= max_depth_reached <= 3:
            return 1.0  # Good - within expected range and under limit
        elif max_depth_reached < max_depth_expected:
            return 0.8  # Partial credit - didn't reach expected depth but under limit
        else:
            return 0.0  # Failed - exceeded limit
