"""StateMixin — session state persistence and versioning."""
from __future__ import annotations
from typing import Any

from logos_adk.state import AgentState


class StateMixin:
    def _persist_state(self, state: AgentState) -> AgentState:
        stored_state = self._resolve_state_store().save(self.name, state)
        self._sessions[state.session_id] = stored_state
        self._sessions.move_to_end(state.session_id)
        self._evict_cold_sessions()
        return stored_state

    def _evict_cold_sessions(self) -> None:
        while len(self._sessions) > self._max_cached_sessions:
            self._sessions.popitem(last=False)

    def list_state_versions(
        self,
        session_id: str | None = None,
        *,
        limit: int = 50,
    ) -> list[dict[str, Any]]:
        """List persisted versions for a session."""
        return self._resolve_state_store().history(self.name, session_id, limit=limit)

    def export_state(
        self,
        session_id: str | None = None,
        *,
        version: int | None = None,
    ) -> dict[str, Any] | None:
        """Export the latest or specified state snapshot."""
        return self._resolve_state_store().export_state(self.name, session_id, version=version)

    def import_state(
        self,
        payload: dict[str, Any],
        *,
        session_id: str | None = None,
    ) -> AgentState:
        """Import a state snapshot and persist it as a new version."""
        state = self._resolve_state_store().import_state(
            self.name,
            payload,
            session_id=session_id,
        )
        self._sessions[state.session_id] = state
        self._sessions.move_to_end(state.session_id)
        self._evict_cold_sessions()
        return state

    def rollback_state(self, session_id: str | None, *, version: int) -> AgentState:
        """Rollback session state to a previous version."""
        state = self._resolve_state_store().rollback(
            self.name,
            session_id,
            version=version,
        )
        self._sessions[session_id] = state
        self._sessions.move_to_end(session_id)
        self._evict_cold_sessions()
        return state

    def state_store_info(self) -> dict[str, Any]:
        """Describe the configured state backend."""
        return self._resolve_state_store().describe()
