"""Storage abstraction for persona debate sessions."""

from __future__ import annotations

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

from .types import (
    Persona,
    PersonaAgreement,
    PersonaCritique,
    PersonaDebateSession,
    PersonaPosition,
)

if TYPE_CHECKING:
    from .._shared.metrics import UsageMetrics


@runtime_checkable
class PersonaDebateStorageProtocol(Protocol):
    """Protocol for persona debate storage implementations.

    Any class that has `session`, `personas`, `positions`, `critiques`, and `agreements`
    properties can be used as storage for the persona debate toolset.

    Example:
        ```python
        class MyCustomStorage:
            def __init__(self):
                self._session: PersonaDebateSession | None = None
                self._personas: dict[str, Persona] = {}
                self._positions: dict[str, PersonaPosition] = {}
                self._critiques: dict[str, PersonaCritique] = {}
                self._agreements: dict[str, PersonaAgreement] = {}

            @property
            def session(self) -> PersonaDebateSession | None:
                return self._session

            @session.setter
            def session(self, value: PersonaDebateSession) -> None:
                self._session = value

            @property
            def personas(self) -> dict[str, Persona]:
                return self._personas

            @personas.setter
            def personas(self, value: Persona) -> None:
                self._personas[value.persona_id] = value

            @property
            def positions(self) -> dict[str, PersonaPosition]:
                return self._positions

            @positions.setter
            def positions(self, value: PersonaPosition) -> None:
                self._positions[value.position_id] = value

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

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

            @property
            def agreements(self) -> dict[str, PersonaAgreement]:
                return self._agreements

            @agreements.setter
            def agreements(self, value: PersonaAgreement) -> None:
                self._agreements[value.agreement_id] = value
        ```
    """

    @property
    def session(self) -> PersonaDebateSession | None:
        """Get the current persona debate session."""
        ...

    @session.setter
    def session(self, value: PersonaDebateSession) -> None:
        """Set the persona debate session."""
        ...

    @property
    def personas(self) -> dict[str, Persona]:
        """Get all personas (persona_id -> Persona)."""
        ...

    @personas.setter
    def personas(self, value: Persona) -> None:
        """Add or update a persona in the dictionary."""
        ...

    @property
    def positions(self) -> dict[str, PersonaPosition]:
        """Get all positions (position_id -> PersonaPosition)."""
        ...

    @positions.setter
    def positions(self, value: PersonaPosition) -> None:
        """Add or update a position in the dictionary."""
        ...

    @property
    def critiques(self) -> dict[str, PersonaCritique]:
        """Get all critiques (critique_id -> PersonaCritique)."""
        ...

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

    @property
    def agreements(self) -> dict[str, PersonaAgreement]:
        """Get all agreements (agreement_id -> PersonaAgreement)."""
        ...

    @agreements.setter
    def agreements(self, value: PersonaAgreement) -> None:
        """Add or update an agreement in the dictionary."""
        ...


@dataclass
class PersonaDebateStorage:
    """Default in-memory persona debate storage.

    Simple implementation that stores persona debate sessions, personas, positions,
    critiques, and agreements in memory. Use this for standalone agents or testing.

    Example:
        ```python
        from pydantic_ai_toolsets import create_persona_debate_toolset, PersonaDebateStorage

        storage = PersonaDebateStorage()
        toolset = create_persona_debate_toolset(storage=storage)

        # After agent runs, access debate state directly
        print(storage.session)
        print(storage.personas)
        print(storage.positions)
        print(storage.critiques)
        print(storage.agreements)

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

    _session: PersonaDebateSession | None = None
    _personas: dict[str, Persona] = field(default_factory=dict)
    _positions: dict[str, PersonaPosition] = field(default_factory=dict)
    _critiques: dict[str, PersonaCritique] = field(default_factory=dict)
    _agreements: dict[str, PersonaAgreement] = 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._session = None
        self._personas = {}
        self._positions = {}
        self._critiques = {}
        self._agreements = {}
        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 session(self) -> PersonaDebateSession | None:
        """Get the current persona debate session."""
        return self._session

    @session.setter
    def session(self, value: PersonaDebateSession) -> None:
        """Set the persona debate session."""
        self._session = value

    @property
    def personas(self) -> dict[str, Persona]:
        """Get all personas (persona_id -> Persona)."""
        return self._personas

    @personas.setter
    def personas(self, value: Persona) -> None:
        """Add or update a persona in the dictionary."""
        self._personas[value.persona_id] = value

    @property
    def positions(self) -> dict[str, PersonaPosition]:
        """Get all positions (position_id -> PersonaPosition)."""
        return self._positions

    @positions.setter
    def positions(self, value: PersonaPosition) -> None:
        """Add or update a position in the dictionary."""
        self._positions[value.position_id] = value

    @property
    def critiques(self) -> dict[str, PersonaCritique]:
        """Get all critiques (critique_id -> PersonaCritique)."""
        return self._critiques

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

    @property
    def agreements(self) -> dict[str, PersonaAgreement]:
        """Get all agreements (agreement_id -> PersonaAgreement)."""
        return self._agreements

    @agreements.setter
    def agreements(self, value: PersonaAgreement) -> None:
        """Add or update an agreement in the dictionary."""
        self._agreements[value.agreement_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 persona debate operations.

        Returns:
            Dictionary with debate counts and metrics.
        """
        total_personas = len(self._personas)
        total_positions = len(self._positions)
        total_critiques = len(self._critiques)
        total_agreements = len(self._agreements)
        current_round = self._session.current_round if self._session else 0
        max_round = self._session.max_rounds if self._session else 0

        return {
            "has_session": 1 if self._session else 0,
            "total_personas": total_personas,
            "total_positions": total_positions,
            "total_critiques": total_critiques,
            "total_agreements": total_agreements,
            "current_round": current_round,
            "max_rounds": max_round,
        }

    def clear(self) -> None:
        """Clear all debate data and reset metrics."""
        self._session = None
        self._personas.clear()
        self._positions.clear()
        self._critiques.clear()
        self._agreements.clear()
        if self._metrics:
            self._metrics.clear()
