import pickle
from typing import Any, Callable, Optional, cast

from redis.asyncio import Redis

from kaya.session import Session, SessionStore


class RedisSessionStore(SessionStore):
    """Redis-backed session store.

    Suitable for multi-process and multi-host deployments: session data is
    shared between all application instances connected to the same Redis
    server.

    Server-side expiry is delegated to Redis key TTLs. When ``max_age`` is
    provided, active sessions slide the expiry window on each access.

    Session data is serialized with ``pickle`` by default, so arbitrary
    Python objects can be stored. Custom serializers can be plugged in via
    the ``dumps``/``loads`` arguments. Only connect this store to a trusted
    Redis server, as pickle deserialization of untrusted data is unsafe.
    """

    def __init__(
        self,
        client: Redis,
        prefix: str = "kaya:session:",
        dumps: Callable[[dict[str, Any]], bytes] = pickle.dumps,
        loads: Callable[[bytes], dict[str, Any]] = pickle.loads,
    ) -> None:
        self._client = client
        self._prefix = prefix
        self._dumps = dumps
        self._loads = loads

    def _key(self, session_id: str) -> str:
        return f"{self._prefix}{session_id}"

    async def load(self, session_id: str, max_age: Optional[int] = None) -> Optional[Session]:
        key = self._key(session_id)
        payload = await self._client.get(key)
        if payload is None:
            return None

        if max_age is not None:
            await self._client.expire(key, max_age)

        data = self._loads(cast(bytes, payload))
        return Session(session_id, data)

    async def save(self, session_id: str, session: Session, max_age: Optional[int] = None) -> None:
        payload = self._dumps(dict(session))
        await self._client.set(self._key(session_id), payload, ex=max_age)

    async def delete(self, session_id: str) -> None:
        await self._client.delete(self._key(session_id))