"""
Redis backend for st-page-state.

Provides optional server-side persistence for PageState classes.
State is loaded from Redis at the beginning of each Streamlit script run
and saved back on exit via the :func:`redis_session` context manager.

Quick start::

    from st_page_state import configure_redis, redis_session

    configure_redis(host="localhost", default_ttl=3600)

    with redis_session():
        ...  # page code

Per-class TTL override::

    class MyState(PageState):
        count: int = StateVar(default=0)

        class Config:
            ttl = 600  # overrides global default_ttl
"""

from __future__ import annotations

import logging
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

@dataclass
class RedisConfig:
    """Connection parameters and global defaults for the Redis backend."""

    host: str = "localhost"
    port: int = 6379
    db: int = 0
    password: Optional[str] = None
    default_ttl: Optional[int] = None
    key_prefix: str = "st_page_state"
    ssl: bool = False
    socket_timeout: Optional[float] = None
    extra_kwargs: Dict[str, Any] = field(default_factory=dict)


# Module-level config singleton – set by configure_redis()
_config: Optional[RedisConfig] = None


def configure_redis(
    host: str = "localhost",
    port: int = 6379,
    db: int = 0,
    password: Optional[str] = None,
    default_ttl: Optional[int] = None,
    key_prefix: str = "st_page_state",
    ssl: bool = False,
    socket_timeout: Optional[float] = None,
    **kwargs: Any,
) -> None:
    """Set the global Redis connection. Call once before any :func:`redis_session` usage.

    *default_ttl* applies to every PageState class unless overridden by ``Config.ttl``.
    Extra *kwargs* are forwarded to ``redis.StrictRedis``.
    """
    global _config
    _config = RedisConfig(
        host=host,
        port=port,
        db=db,
        password=password,
        default_ttl=default_ttl,
        key_prefix=key_prefix,
        ssl=ssl,
        socket_timeout=socket_timeout,
        extra_kwargs=kwargs,
    )
    logger.debug("st-page-state: Redis configured – %s:%d db=%d", host, port, db)


# ---------------------------------------------------------------------------
# Client
# ---------------------------------------------------------------------------

class RedisClient:
    """Thin wrapper around ``redis.StrictRedis``.

    Owns the TCP connection, key naming, and TTL resolution.
    One instance is created per :func:`redis_session` invocation.
    """

    def __init__(self, config: RedisConfig) -> None:
        self._config = config
        self._client = self._connect(config)

    @staticmethod
    def _connect(config: RedisConfig):
        try:
            import redis
            
        except ImportError as exc:
            raise ImportError(
                "The 'redis' package is required for Redis state persistence.\n"
                "Install it with:  pip install redis"
            ) from exc

        return redis.StrictRedis(
            host=config.host,
            port=config.port,
            db=config.db,
            password=config.password,
            ssl=config.ssl,
            socket_timeout=config.socket_timeout,
            decode_responses=True,
            **config.extra_kwargs,
        )

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

    def resolve_ttl(self, state_cls) -> Optional[int]:
        """Return effective TTL: ``Config.ttl`` on the class → global ``default_ttl`` → ``None``."""
        user_config = getattr(state_cls, "Config", None)
        if user_config is not None and hasattr(user_config, "ttl"):
            return getattr(user_config, "ttl")
        return self._config.default_ttl

    def load(self, class_name: str, session_id: str) -> Optional[Dict[str, Any]]:
        """Load a stored state dict, or return ``None`` if absent."""
        from ..utils.converters import deserialize_state

        key = self._key(class_name, session_id)
        try:
            raw = self._client.get(key)
            return deserialize_state(raw) if raw is not None else None
        except Exception as exc:
            logger.warning("st-page-state: Redis load failed [%s]: %s", key, exc)
            return None

    def save(self, class_name: str, session_id: str, data: Dict[str, Any], ttl: Optional[int]) -> None:
        """Persist *data* to Redis, optionally with a TTL (seconds)."""
        from ..utils.converters import serialize_state

        key = self._key(class_name, session_id)
        try:
            payload = serialize_state(data)
            if ttl:
                self._client.setex(key, ttl, payload)
            else:
                self._client.set(key, payload)
        except Exception as exc:
            logger.warning("st-page-state: Redis save failed [%s]: %s", key, exc)

    def delete(self, class_name: str, session_id: str) -> None:
        """Remove the key for *class_name* / *session_id*."""
        key = self._key(class_name, session_id)
        try:
            self._client.delete(key)
        except Exception as exc:
            logger.warning("st-page-state: Redis delete failed [%s]: %s", key, exc)

    def ping(self) -> bool:
        """Return ``True`` if Redis is reachable."""
        try:
            return bool(self._client.ping())
        except Exception:
            return False


# ---------------------------------------------------------------------------
# Session helpers
# ---------------------------------------------------------------------------

def _session_id() -> str:
    """Return the current Streamlit session ID (falls back to ``"default"``)."""
    try:
        from streamlit.runtime.scriptrunner import get_script_run_ctx
        ctx = get_script_run_ctx()
        if ctx is not None:
            return ctx.session_id
    except Exception:
        pass
    return "default"


# ---------------------------------------------------------------------------
# Context manager
# ---------------------------------------------------------------------------

@contextmanager
def redis_session():
    """Load state from Redis on entry, save it back on exit.

    If :func:`configure_redis` was never called, this is a transparent no-op.

    Each registered PageState class is serialized independently.
    TTL is resolved per class (``Config.ttl`` → global ``default_ttl`` → no expiry).

    .. note::

       Streamlit has no session-end hook. Redis key TTL is the primary
       mechanism for cleaning up stale state. For finer-grained idle
       detection, use ``st.fragment(run_every=...)`` in your app code.
    """
    if _config is None:
        yield
        return

    import streamlit as st
    from ..core.meta import _PAGE_STATE_REGISTRY, SESSION_STATE_KEY

    client = RedisClient(_config)
    sid = _session_id()

    # ── LOAD ────────────────────────────────────────────────────────────────
    for class_name, state_cls in _PAGE_STATE_REGISTRY.items():
        stored = client.load(class_name, sid)
        if not stored:
            continue

        st.session_state.setdefault(SESSION_STATE_KEY, {}).setdefault(class_name, {})
        class_ns = st.session_state[SESSION_STATE_KEY][class_name]

        for field_name, value in stored.items():
            if field_name in state_cls._model_metadata:
                class_ns[field_name] = value

        logger.debug(
            "st-page-state: Loaded %d field(s) for '%s' (session=%s)",
            len(stored), class_name, sid,
        )

    try:
        yield

    finally:
        # ── SAVE ────────────────────────────────────────────────────────────
        for class_name, state_cls in _PAGE_STATE_REGISTRY.items():
            class_ns = st.session_state.get(SESSION_STATE_KEY, {}).get(class_name)
            if not class_ns:
                continue

            ttl = client.resolve_ttl(state_cls)
            client.save(class_name, sid, dict(class_ns), ttl)
            logger.debug(
                "st-page-state: Saved %d field(s) for '%s' (session=%s, ttl=%s)",
                len(class_ns), class_name, sid, ttl,
            )
