"""Cascade State Store — Persistent cascade detection state with Redis INCR backend.

Provides CascadeStateStore Protocol with InMemoryCascadeStateStore and
RedisCascadeStateStore backends. Enables CascadeDetector state to survive
restarts and work across multiple instances.

Key schema (Redis):
    agentkit:cascade:interactions:{session_id}  → INCR counter with TTL
    agentkit:cascade:depths:{session_id}        → SET counter with TTL
"""

import logging
import time
from typing import Protocol, runtime_checkable

# redis 可选依赖；未安装时回退为 Exception 以保留原 catch-all 语义（降级到 fallback）
try:
    from redis.exceptions import RedisError as _RedisError
except ImportError:
    _RedisError = Exception

logger = logging.getLogger(__name__)


@runtime_checkable
class CascadeStateStore(Protocol):
    """Persistent cascade detection state interface."""

    def increment_interaction(self, session_id: str) -> int:
        """Atomically increment interaction count. Returns new count."""
        ...

    def get_interaction(self, session_id: str) -> int:
        """Get current interaction count."""
        ...

    def set_depth(self, session_id: str, depth: int) -> None:
        """Set loop depth for a session."""
        ...

    def get_depth(self, session_id: str) -> int:
        """Get current loop depth."""
        ...

    def reset(self, session_id: str) -> None:
        """Reset all counters for a session."""
        ...


# ---------------------------------------------------------------------------
# InMemoryCascadeStateStore
# ---------------------------------------------------------------------------


class InMemoryCascadeStateStore:
    """In-memory cascade state store (default, process-local).

    Supports optional session TTL to prevent unbounded memory growth.
    Expired entries are lazily cleaned up on access.
    """

    DEFAULT_SESSION_TTL = 86400  # 24 hours

    def __init__(self, session_ttl: int = 86400):
        self._session_ttl = session_ttl
        self._interaction_counts: dict[str, int] = {}
        self._loop_depths: dict[str, int] = {}
        self._timestamps: dict[str, float] = {}

    def _is_expired(self, session_id: str) -> bool:
        ts = self._timestamps.get(session_id)
        if ts is None:
            return False
        return (time.monotonic() - ts) > self._session_ttl

    def _cleanup_expired(self) -> None:
        """Lazy cleanup: remove expired sessions."""
        expired = [sid for sid in self._timestamps if self._is_expired(sid)]
        for sid in expired:
            self._interaction_counts.pop(sid, None)
            self._loop_depths.pop(sid, None)
            self._timestamps.pop(sid, None)

    def _touch(self, session_id: str) -> None:
        self._timestamps[session_id] = time.monotonic()

    def increment_interaction(self, session_id: str) -> int:
        self._cleanup_expired()
        self._interaction_counts[session_id] = self._interaction_counts.get(session_id, 0) + 1
        self._touch(session_id)
        return self._interaction_counts[session_id]

    def get_interaction(self, session_id: str) -> int:
        if self._is_expired(session_id):
            self.reset(session_id)
            return 0
        return self._interaction_counts.get(session_id, 0)

    def set_depth(self, session_id: str, depth: int) -> None:
        self._touch(session_id)
        self._loop_depths[session_id] = depth

    def get_depth(self, session_id: str) -> int:
        if self._is_expired(session_id):
            self.reset(session_id)
            return 0
        return self._loop_depths.get(session_id, 0)

    def reset(self, session_id: str) -> None:
        self._interaction_counts.pop(session_id, None)
        self._loop_depths.pop(session_id, None)
        self._timestamps.pop(session_id, None)


# ---------------------------------------------------------------------------
# RedisCascadeStateStore
# ---------------------------------------------------------------------------


class RedisCascadeStateStore:
    """Redis-backed cascade state store using INCR for atomic increments.

    Key schema:
        agentkit:cascade:interactions:{session_id}  → INCR counter with TTL
        agentkit:cascade:depths:{session_id}         → SET counter with TTL
    """

    INTER_PREFIX = "agentkit:cascade:interactions:"
    DEPTH_PREFIX = "agentkit:cascade:depths:"
    SESSION_TTL = 86400  # 24 hours — sessions rarely last longer

    def __init__(self, redis_url: str = "redis://localhost:6379", session_ttl: int = 86400):
        self._redis_url = redis_url
        self._session_ttl = session_ttl
        self._sync_redis: object = None
        self._fallback: InMemoryCascadeStateStore | None = None
        self._degraded = False

    def _get_sync_redis(self):
        """Get or create a persistent sync Redis client (connection pool backed)."""
        if self._sync_redis is None:
            import redis as sync_redis
            self._sync_redis = sync_redis.from_url(
                self._redis_url, decode_responses=True
            )
        return self._sync_redis

    def _degrade_to_fallback(self) -> None:
        if not self._degraded:
            self._degraded = True
            if self._fallback is None:
                self._fallback = InMemoryCascadeStateStore(session_ttl=self._session_ttl)
            logger.warning("Redis cascade store unreachable, degraded to in-memory")

    def increment_interaction(self, session_id: str) -> int:
        if self._degraded and self._fallback is not None:
            return self._fallback.increment_interaction(session_id)
        try:
            r = self._get_sync_redis()
            key = f"{self.INTER_PREFIX}{session_id}"
            pipe = r.pipeline()
            pipe.incr(key)
            pipe.expire(key, self._session_ttl)
            results = pipe.execute()
            return results[0]
        except (ImportError, OSError, _RedisError, ValueError, KeyError, RuntimeError, TypeError) as e:
            logger.warning(f"Redis cascade increment failed: {e}")
            self._degrade_to_fallback()
            if self._fallback is not None:
                return self._fallback.increment_interaction(session_id)
            return 0

    def get_interaction(self, session_id: str) -> int:
        if self._degraded and self._fallback is not None:
            return self._fallback.get_interaction(session_id)
        try:
            r = self._get_sync_redis()
            val = r.get(f"{self.INTER_PREFIX}{session_id}")
            return int(val) if val is not None else 0
        except (ImportError, OSError, _RedisError, ValueError, KeyError, RuntimeError, TypeError) as e:
            logger.warning(f"Redis cascade get failed: {e}")
            if self._fallback is not None:
                return self._fallback.get_interaction(session_id)
            return 0

    def set_depth(self, session_id: str, depth: int) -> None:
        if self._degraded and self._fallback is not None:
            self._fallback.set_depth(session_id, depth)
            return
        try:
            r = self._get_sync_redis()
            key = f"{self.DEPTH_PREFIX}{session_id}"
            pipe = r.pipeline()
            pipe.set(key, depth)
            pipe.expire(key, self._session_ttl)
            pipe.execute()
        except (ImportError, OSError, _RedisError, ValueError, KeyError, RuntimeError, TypeError) as e:
            logger.warning(f"Redis cascade set_depth failed: {e}")
            self._degrade_to_fallback()
            if self._fallback is not None:
                self._fallback.set_depth(session_id, depth)

    def get_depth(self, session_id: str) -> int:
        if self._degraded and self._fallback is not None:
            return self._fallback.get_depth(session_id)
        try:
            r = self._get_sync_redis()
            val = r.get(f"{self.DEPTH_PREFIX}{session_id}")
            return int(val) if val is not None else 0
        except (ImportError, OSError, _RedisError, ValueError, KeyError, RuntimeError, TypeError) as e:
            logger.warning(f"Redis cascade get_depth failed: {e}")
            if self._fallback is not None:
                return self._fallback.get_depth(session_id)
            return 0

    def reset(self, session_id: str) -> None:
        if self._degraded and self._fallback is not None:
            self._fallback.reset(session_id)
            return
        try:
            r = self._get_sync_redis()
            pipe = r.pipeline()
            pipe.delete(f"{self.INTER_PREFIX}{session_id}")
            pipe.delete(f"{self.DEPTH_PREFIX}{session_id}")
            pipe.execute()
        except (ImportError, OSError, _RedisError, ValueError, KeyError, RuntimeError, TypeError) as e:
            logger.warning(f"Redis cascade reset failed: {e}")
            self._degrade_to_fallback()
            if self._fallback is not None:
                self._fallback.reset(session_id)

    def close(self) -> None:
        """Close the Redis connection pool."""
        if self._sync_redis is not None:
            self._sync_redis.close()
            self._sync_redis = None


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------


def create_cascade_state_store(
    backend: str = "auto",
    redis_url: str = "redis://localhost:6379",
    session_ttl: int = 86400,
) -> CascadeStateStore:
    """Create a cascade state store backend."""
    if backend in ("auto", "redis"):
        try:
            import redis  # noqa: F401
            return RedisCascadeStateStore(redis_url=redis_url, session_ttl=session_ttl)
        except ImportError:
            logger.warning("redis package not available, falling back to in-memory cascade store")
    return InMemoryCascadeStateStore(session_ttl=session_ttl)
