"""Storage state evaluator."""

from __future__ import annotations

from typing import Any

from pydantic_ai_toolsets.evals.base import EvaluationResult


class StorageStateEvaluator:
    """Evaluates storage state: structure, content quality, completeness."""

    def evaluate_storage_present(self, result: EvaluationResult) -> float:
        """Check if storage state is present.

        Args:
            result: Evaluation result.

        Returns:
            1.0 if storage state exists, 0.0 otherwise.
        """
        return 1.0 if result.storage_state else 0.0

    def evaluate_storage_structure(
        self, result: EvaluationResult, expected_keys: list[str] | None = None
    ) -> float:
        """Evaluate if storage has expected structure.

        Args:
            result: Evaluation result.
            expected_keys: Expected keys in storage state.

        Returns:
            Score from 0.0 to 1.0.
        """
        if not result.storage_state:
            return 0.0

        if not expected_keys:
            # If no expected keys, just check that storage has some structure
            return 1.0 if result.storage_state else 0.0

        present_keys = set(result.storage_state.keys())
        expected_set = set(expected_keys)

        if expected_set.issubset(present_keys):
            return 1.0

        # Partial credit
        overlap = len(expected_set & present_keys)
        return overlap / len(expected_set) if expected_set else 1.0

    def evaluate_storage_content_quality(
        self, result: EvaluationResult, min_items: int = 1
    ) -> float:
        """Evaluate quality of storage content.

        Args:
            result: Evaluation result.
            min_items: Minimum number of items expected in storage.

        Returns:
            Score from 0.0 to 1.0.
        """
        if not result.storage_state:
            return 0.0

        # Count items in various storage types
        total_items = 0

        # Count todos
        if "todos" in result.storage_state:
            total_items += len(result.storage_state["todos"])

        # Count thoughts
        if "thoughts" in result.storage_state:
            total_items += len(result.storage_state["thoughts"])

        # Count nodes
        if "nodes" in result.storage_state:
            total_items += len(result.storage_state["nodes"])

        # Count candidates
        if "candidates" in result.storage_state:
            total_items += len(result.storage_state["candidates"])

        # Count outputs
        if "outputs" in result.storage_state:
            total_items += len(result.storage_state["outputs"])

        # Score based on whether minimum items are present
        if total_items >= min_items:
            return 1.0
        elif min_items > 0:
            return total_items / min_items
        else:
            return 1.0 if total_items > 0 else 0.0

    def evaluate_storage_completeness(
        self, result: EvaluationResult, expected_counts: dict[str, int] | None = None
    ) -> float:
        """Evaluate completeness of storage.

        Args:
            result: Evaluation result.
            expected_counts: Expected counts for different storage types.

        Returns:
            Score from 0.0 to 1.0.
        """
        if not result.storage_state:
            return 0.0

        if not expected_counts:
            return 1.0

        scores = []
        for key, expected_count in expected_counts.items():
            actual_count = 0
            if key in result.storage_state:
                value = result.storage_state[key]
                if isinstance(value, list):
                    actual_count = len(value)
                elif isinstance(value, dict):
                    actual_count = len(value)
                elif isinstance(value, (int, float)):
                    actual_count = int(value)

            if expected_count > 0:
                score = min(actual_count / expected_count, 1.0)
            else:
                score = 1.0 if actual_count == 0 else 0.0

            scores.append(score)

        return sum(scores) / len(scores) if scores else 1.0

    def evaluate_all(
        self,
        result: EvaluationResult,
        expected_keys: list[str] | None = None,
        min_items: int = 1,
        expected_counts: dict[str, int] | None = None,
    ) -> dict[str, float]:
        """Evaluate all storage state metrics.

        Args:
            result: Evaluation result.
            expected_keys: Expected storage keys.
            min_items: Minimum items expected.
            expected_counts: Expected counts per storage type.

        Returns:
            Dictionary with storage state scores.
        """
        return {
            "storage_present": self.evaluate_storage_present(result),
            "structure_correctness": self.evaluate_storage_structure(
                result, expected_keys
            ),
            "content_quality": self.evaluate_storage_content_quality(
                result, min_items
            ),
            "completeness": self.evaluate_storage_completeness(
                result, expected_counts
            ),
        }

