"""Storage abstraction for reflection."""

from __future__ import annotations

import sys
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Protocol, runtime_checkable

from .types import Critique, ReflectionOutput

if TYPE_CHECKING:
    from .._shared.metrics import UsageMetrics


@runtime_checkable
class ReflectionStorageProtocol(Protocol):
    """Protocol for reflection storage implementations.

    Any class that has `outputs` and `critiques` properties can be used
    as storage for the reflection toolset.

    Example:
        ```python
        class MyCustomStorage:
            def __init__(self):
                self._outputs: dict[str, ReflectionOutput] = {}
                self._critiques: dict[str, Critique] = {}

            @property
            def outputs(self) -> dict[str, ReflectionOutput]:
                return self._outputs

            @outputs.setter
            def outputs(self, value: ReflectionOutput) -> None:
                self._outputs[value.output_id] = value

            @property
            def critiques(self) -> dict[str, Critique]:
                return self._critiques

            @critiques.setter
            def critiques(self, value: Critique) -> None:
                self._critiques[value.critique_id] = value
        ```
    """

    @property
    def outputs(self) -> dict[str, ReflectionOutput]:
        """Get the current dictionary of outputs (output_id -> ReflectionOutput)."""
        ...

    @outputs.setter
    def outputs(self, value: ReflectionOutput) -> None:
        """Add or update an output in the dictionary."""
        ...

    @property
    def critiques(self) -> dict[str, Critique]:
        """Get the current dictionary of critiques (critique_id -> Critique)."""
        ...

    @critiques.setter
    def critiques(self, value: Critique) -> None:
        """Add or update a critique in the dictionary."""
        ...


@dataclass
class ReflectionStorage:
    """Default in-memory reflection storage.

    Simple implementation that stores outputs and critiques in memory.
    Use this for standalone agents or testing.

    Example:
        ```python
        from pydantic_ai_toolsets import create_reflection_toolset, ReflectionStorage

        storage = ReflectionStorage()
        toolset = create_reflection_toolset(storage=storage)

        # After agent runs, access outputs and critiques directly
        print(storage.outputs)
        print(storage.critiques)

        # With metrics tracking
        storage = ReflectionStorage(track_usage=True)
        toolset = create_reflection_toolset(storage=storage)
        print(storage.metrics.total_tokens())
        ```
    """

    _outputs: dict[str, ReflectionOutput] = field(default_factory=dict)
    _critiques: dict[str, Critique] = field(default_factory=dict)
    _metrics: UsageMetrics | None = field(default=None)

    def __init__(self, *, track_usage: bool = False) -> None:
        """Initialize storage with optional metrics tracking.

        Args:
            track_usage: If True, enables usage metrics collection.
        """
        self._outputs = {}
        self._critiques = {}
        self._metrics = None
        if track_usage:
            import os

            toolsets_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            if toolsets_dir not in sys.path:
                sys.path.insert(0, toolsets_dir)
            from .._shared.metrics import UsageMetrics

            self._metrics = UsageMetrics()

    @property
    def outputs(self) -> dict[str, ReflectionOutput]:
        """Get the current dictionary of outputs."""
        return self._outputs

    @outputs.setter
    def outputs(self, value: ReflectionOutput) -> None:
        """Add or update an output in the dictionary."""
        self._outputs[value.output_id] = value

    @property
    def critiques(self) -> dict[str, Critique]:
        """Get the current dictionary of critiques."""
        return self._critiques

    @critiques.setter
    def critiques(self, value: Critique) -> None:
        """Add or update a critique in the dictionary."""
        self._critiques[value.critique_id] = value

    @property
    def metrics(self) -> UsageMetrics | None:
        """Get usage metrics if tracking is enabled."""
        return self._metrics

    def get_statistics(self) -> dict[str, int | float]:
        """Get summary statistics about reflection operations.

        Returns:
            Dictionary with output and critique counts.
        """
        total_outputs = len(self._outputs)
        final_outputs = sum(1 for o in self._outputs.values() if o.is_final)
        max_cycle = max((o.cycle for o in self._outputs.values()), default=0)
        avg_quality = None
        quality_scores = [o.quality_score for o in self._outputs.values() if o.quality_score is not None]
        if quality_scores:
            avg_quality = sum(quality_scores) / len(quality_scores)

        stats: dict[str, int | float] = {
            "total_outputs": total_outputs,
            "final_outputs": final_outputs,
            "max_cycle": max_cycle,
            "total_critiques": len(self._critiques),
        }
        if avg_quality is not None:
            stats["avg_quality_score"] = avg_quality

        return stats

    def clear(self) -> None:
        """Clear all outputs, critiques, and reset metrics."""
        self._outputs.clear()
        self._critiques.clear()
        if self._metrics:
            self._metrics.clear()
