state.py

# src/bot/handlers/state.py
from dataclasses import dataclass
from typing import Optional
import asyncio
import json
from pathlib import Path
import os

RUNTIME_DIR = Path(os.getenv('RUNTIME_DIR', 'runtime'))
RUNTIME_DIR.mkdir(parents=True, exist_ok=True)


@dataclass
class QtyState:
    product: str
    msg_id: int


class PendingQty:
    def __init__(self):
        self._store = {}
        self._lock = asyncio.Lock()

    async def set(self, user_id: int, product: str, msg_id: int):
        async with self._lock:
            self._store[user_id] = QtyState(product, msg_id)

    async def get(self, user_id: int) -> Optional[QtyState]:
        async with self._lock:
            return self._store.get(user_id)

    async def clear(self, user_id: int):
        async with self._lock:
            self._store.pop(user_id, None)


@dataclass
class PremiumState:
    plan: str


class PendingPremium:
    """In-memory pending premium selection with optional disk persistence.

    Persists each user's pending plan to: runtime/pending_premium_{user_id}.json
    Uses run_in_executor to avoid blocking the event loop on file I/O.
    """
    def __init__(self):
        self._store: dict[int, PremiumState] = {}
        self._lock = asyncio.Lock()

    def _path_for(self, user_id: int) -> Path:
        return RUNTIME_DIR / f'pending_premium_{user_id}.json'

    async def _write_file(self, path: Path, data: dict):
        loop = asyncio.get_event_loop()

        def _write():
            with open(path, 'w', encoding='utf-8') as f:
                json.dump(data, f)

        await loop.run_in_executor(None, _write)

    async def _read_file(self, path: Path) -> Optional[dict]:
        loop = asyncio.get_event_loop()

        def _read():
            if not path.exists():
                return None
            with open(path, 'r', encoding='utf-8') as f:
                return json.load(f)

        return await loop.run_in_executor(None, _read)

    async def _remove_file(self, path: Path):
        loop = asyncio.get_event_loop()

        def _rm():
            try:
                if path.exists():
                    path.unlink()
            except Exception:
                pass

        await loop.run_in_executor(None, _rm)

    async def set(self, user_id: int, plan: str):
        """Set pending plan for user and persist to disk."""
        async with self._lock:
            self._store[user_id] = PremiumState(plan)
            data = {'plan': plan}
            path = self._path_for(user_id)
            try:
                await self._write_file(path, data)
            except Exception:
                # best-effort: don't raise to calling handlers
                pass

    async def get(self, user_id: int) -> Optional[PremiumState]:
        async with self._lock:
            st = self._store.get(user_id)
            if st:
                return st
        # if not in memory, try to read from disk
        path = self._path_for(user_id)
        try:
            data = await self._read_file(path)
            if data and 'plan' in data:
                ps = PremiumState(data['plan'])
                # cache in memory for future
                async with self._lock:
                    self._store[user_id] = ps
                return ps
        except Exception:
            return None
        return None

    async def clear(self, user_id: int):
        async with self._lock:
            self._store.pop(user_id, None)
        path = self._path_for(user_id)
        try:
            await self._remove_file(path)
        except Exception:
            pass