memory/session_store.py

# memory/session_store.py
import uuid
import datetime
import logging
from typing import Dict, Any, Optional

logger = logging.getLogger(__name__)


class SessionStore:
    """
    Хранилище состояния сессии.
    Полностью поддерживает логику: 1/0/2 меню, Custom Plot Mode и счетчики попыток.
    """

    def __init__(self):
        self._sessions: Dict[str, Dict[str, Any]] = {}

    def create_session(self) -> str:
        """Создает новую сессию с полным набором полей"""
        session_id = str(uuid.uuid4())
        self._sessions[session_id] = {
            "session_id": session_id,
            "created_at": datetime.datetime.utcnow().isoformat(),

            # --- Логика машины состояний ---
            "step": "initial",
            "attempts": 0,
            "suggested_plot": None,  # Хранит текст предложенной истории
            "custom_plot_mode": False,  # ✅ Флаг строгого режима (критично для вашей логики)

            # --- Данные диалога ---
            "history": [],
            "user_inputs": [],

            # --- Результаты текущего поиска ---
            "original_plot": "",
            "improved_plot": "",
            "movie_overview": "",
            "search_results": [],
            "final_recommendations": []
        }
        return session_id

    def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
        """Возвращает копию состояния сессии"""
        return self._sessions.get(session_id)

    # --- Методы управления состоянием ---

    def add_user_input(self, session_id: str, text: str):
        if session_id in self._sessions:
            self._sessions[session_id]["user_inputs"].append(text)

    def add_history(self, session_id: str, user_msg: str, assistant_msg: str):
        if session_id in self._sessions:
            self._sessions[session_id]["history"].append({"role": "user", "content": user_msg})
            self._sessions[session_id]["history"].append({"role": "assistant", "content": assistant_msg})

    def update_state(self, session_id: str, key: str, value: Any):
        """Универсальное обновление поля"""
        if session_id in self._sessions:
            self._sessions[session_id][key] = value

    def increment_attempts(self, session_id: str):
        if session_id in self._sessions:
            self._sessions[session_id]["attempts"] += 1

    def reset_attempts(self, session_id: str):
        if session_id in self._sessions:
            self._sessions[session_id]["attempts"] = 0

    def set_custom_mode(self, session_id: str, is_active: bool):
        """Управление флагом строгого режима"""
        if session_id in self._sessions:
            self._sessions[session_id]["custom_plot_mode"] = is_active

    def clear_session_data(self, session_id: str):
        """Мягкий сброс (для кнопки Reset)"""
        if session_id in self._sessions:
            base = self._sessions[session_id]
            # Сбрасываем динамические данные
            base["history"] = []
            base["user_inputs"] = []
            base["original_plot"] = ""
            base["final_recommendations"] = []
            base["attempts"] = 0
            base["suggested_plot"] = None
            base["custom_plot_mode"] = False