import streamlit as st
import logging
import threading
from contextlib import contextmanager
from typing import Any, Dict, List, Optional, Tuple

from ..core.meta import _PAGE_STATE_REGISTRY, SESSION_STATE_KEY
from ..utils.converters import deserialize_state, serialize_state

logger = logging.getLogger(__name__)


class RedisBackend:
    """Single entry-point for Redis-backed state persistence.

    Holds the connection configuration, the ``redis.StrictRedis`` client,
    and exposes :meth:`session` as the context-manager wrapper around
    Streamlit's script-run lifecycle.

    Parameters match ``redis.StrictRedis`` plus two extras:

    *default_ttl* — seconds before a key expires (overridable per-class
    via ``Config.ttl``).  ``None`` means no expiry.

    *key_prefix* — namespaces every Redis key as
    ``<prefix>:<session_id>:<ClassName>``.
    """

    def __init__(
        self,
        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:
        
        self.host = host
        self.port = port
        self.db = db
        self.password = password
        self.default_ttl = default_ttl
        self.key_prefix = key_prefix
        self.ssl = ssl
        self.socket_timeout = socket_timeout
        self.extra_kwargs = kwargs

        self._client = self._connect()
        self._last_save_thread: Optional[threading.Thread] = None

        logger.debug(f"Redis configured – {host}:{port} db={db}")

    # -- connection ----------------------------------------------------------

    def _connect(self):
        try:
            import redis as _redis

        except ImportError as exc:
            raise ImportError(
                "The 'redis' package is required for Redis state persistence.\n"
                "Install it with:  pip install st-page-state[redis]"
            ) from exc

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

    # -- key / TTL -----------------------------------------------------------

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

    def resolve_ttl(self, state_cls) -> Optional[int]:
        """``Config.ttl`` on the class → ``default_ttl`` → ``None``."""

        cfg = getattr(state_cls, "Config", None)
        if cfg is not None and hasattr(cfg, "ttl"):
            return getattr(cfg, "ttl")
        
        return self.default_ttl

    # -- CRUD ----------------------------------------------------------------

    def load(self, class_name: str, session_id: str) -> Optional[Dict[str, Any]]:
        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(f"Redis load failed [{key}]: {exc}")
            return None

    def save(self, class_name: str, session_id: str, data: Dict[str, Any], ttl: Optional[int]) -> None:
        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(f"Redis save failed [{key}]: {exc}")

    def delete(self, class_name: str, session_id: str) -> None:
        key = self._key(class_name, session_id)
        try:

            self._client.delete(key)

        except Exception as exc:
            logger.warning(f"Redis delete failed [{key}]: {exc}")

    def ping(self) -> bool:
        try:
            return bool(self._client.ping())
        
        except Exception:
            return False

    # -- session context manager ---------------------------------------------

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

        Each registered PageState class is serialized independently.
        TTL is resolved per class (``Config.ttl`` → ``default_ttl`` → no expiry).
        """
        sid = self._session_id()

        # ---
        # LOAD
        existing = st.session_state.get(SESSION_STATE_KEY, {})

        for class_name, state_cls in _PAGE_STATE_REGISTRY.items():
            if class_name in existing and existing[class_name]:
                continue

            stored = self.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(f"Restored {len(stored)} field(s) for '{class_name}' from Redis (session={sid})")

        try:
            yield
        finally:
            self._save(sid)

    # -- internals -----------------------------------------------------------

    @staticmethod
    def _session_id() -> str:
        """Current Streamlit session ID, or ``"default"`` outside a session."""
        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"

    def _save(self, session_id: str) -> None:
        """Snapshot every class namespace and persist in a daemon thread."""
        payloads: List[Tuple[str, Dict[str, Any], Optional[int]]] = []

        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

            payloads.append((class_name, dict(class_ns), self.resolve_ttl(state_cls)))

        if payloads:

            t = threading.Thread(target=self._save_worker, args=(session_id, payloads), daemon=True)
            t.start()
            self._last_save_thread = t

    def _save_worker(self, session_id: str, payloads: "List[Tuple[str, Dict[str, Any], Optional[int]]]") -> None:
        for class_name, data, ttl in payloads:
            self.save(class_name, session_id, data, ttl)

            logger.debug(f"Saved {len(data)} field(s) for '{class_name}' (session={session_id}, ttl={ttl})")
