"""Base class for thinking/cognition toolset evaluations."""

from typing import Any, Callable

from pydantic_ai.toolsets import FunctionToolset

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


class BaseThinkingEvaluator(ToolsetEvaluator[str, Any]):
    """Base evaluator for thinking/cognition toolsets."""

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

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

        for case in THINKING_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)

            # 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.case_name = case.name
            results.append(result)

        return results

