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

from __future__ import annotations

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

from .types import Persona, PersonaResponse, PersonaSession


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

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

    Example:
        ```python
        class MyCustomStorage:
            def __init__(self):
                self._session: PersonaSession | None = None
                self._personas: dict[str, Persona] = {}
                self._responses: dict[str, PersonaResponse] = {}

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

            @session.setter
            def session(self, value: PersonaSession) -> 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 responses(self) -> dict[str, PersonaResponse]:
                return self._responses

            @responses.setter
            def responses(self, value: PersonaResponse) -> None:
                self._responses[value.response_id] = value
        ```
    """

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

    @session.setter
    def session(self, value: PersonaSession) -> None:
        """Set the persona 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 responses(self) -> dict[str, PersonaResponse]:
        """Get all responses (response_id -> PersonaResponse)."""
        ...

    @responses.setter
    def responses(self, value: PersonaResponse) -> None:
        """Add or update a response in the dictionary."""
        ...


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

    Simple implementation that stores persona sessions, personas, and responses in memory.
    Use this for standalone agents or testing.

    Example:
        ```python
        from pydantic_ai_toolsets import create_persona_toolset, PersonaStorage

        storage = PersonaStorage()
        toolset = create_persona_toolset(storage=storage)

        # After agent runs, access persona state directly
        print(storage.session)
        print(storage.personas)
        print(storage.responses)
        ```
    """

    _session: PersonaSession | None = None
    _personas: dict[str, Persona] = field(default_factory=lambda: {})
    _responses: dict[str, PersonaResponse] = field(default_factory=lambda: {})

    @property
    def session(self) -> PersonaSession | None:
        """Get the current persona session."""
        return self._session

    @session.setter
    def session(self, value: PersonaSession) -> None:
        """Set the persona 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 responses(self) -> dict[str, PersonaResponse]:
        """Get all responses (response_id -> PersonaResponse)."""
        return self._responses

    @responses.setter
    def responses(self, value: PersonaResponse) -> None:
        """Add or update a response in the dictionary."""
        self._responses[value.response_id] = value

