"""Redis backend implementation for state storage."""

import json
from typing import Optional, List
import redis
from redis.exceptions import RedisError

from .base import StateBackend
from ..schemas import UnifiedState


class RedisBackend(StateBackend):
    """Redis-based state storage backend."""
    
    def __init__(
        self,
        host: str = "localhost",
        port: int = 6379,
        db: int = 0,
        password: Optional[str] = None,
        decode_responses: bool = True,
        **kwargs
    ):
        """
        Initialize Redis backend.
        
        Args:
            host: Redis host
            port: Redis port
            db: Redis database number
            password: Redis password (if required)
            decode_responses: Whether to decode responses as strings
            **kwargs: Additional Redis connection parameters
        """
        self.client = redis.Redis(
            host=host,
            port=port,
            db=db,
            password=password,
            decode_responses=decode_responses,
            **kwargs
        )
        # Test connection
        try:
            self.client.ping()
        except RedisError as e:
            raise ConnectionError(f"Failed to connect to Redis: {e}")
    
    def write_state(self, state: UnifiedState) -> bool:
        """Write state to Redis."""
        try:
            key = state.get_key()
            value = state.to_json()
            
            # Write state
            self.client.set(key, value)
            
            # Add to agents list
            agents_key = state.get_agents_list_key()
            self.client.sadd(agents_key, state.agent_id)
            
            # Append to history (keep last 100 entries)
            history_key = state.get_history_key()
            history_entry = {
                "timestamp": state.metadata.updated_at.isoformat(),
                "state": state.state.to_dict(),
                "metadata": state.metadata.to_dict(),
            }
            self.client.lpush(history_key, json.dumps(history_entry))
            self.client.ltrim(history_key, 0, 99)  # Keep last 100 entries
            
            return True
        except RedisError as e:
            print(f"Error writing state to Redis: {e}")
            return False
    
    def read_state(self, session_id: str, agent_id: str) -> Optional[UnifiedState]:
        """Read state from Redis."""
        try:
            key = f"uass:session:{session_id}:agent:{agent_id}:state"
            value = self.client.get(key)
            
            if value is None:
                return None
            
            return UnifiedState.from_json(value)
        except (RedisError, json.JSONDecodeError) as e:
            print(f"Error reading state from Redis: {e}")
            return None
    
    def delete_state(self, session_id: str, agent_id: str) -> bool:
        """Delete state from Redis."""
        try:
            key = f"uass:session:{session_id}:agent:{agent_id}:state"
            history_key = f"uass:session:{session_id}:agent:{agent_id}:history"
            agents_key = f"uass:session:{session_id}:agents"
            
            # Delete state and history
            self.client.delete(key)
            self.client.delete(history_key)
            
            # Remove from agents list
            self.client.srem(agents_key, agent_id)
            
            return True
        except RedisError as e:
            print(f"Error deleting state from Redis: {e}")
            return False
    
    def list_agents(self, session_id: str) -> List[str]:
        """List all agent IDs in a session."""
        try:
            agents_key = f"uass:session:{session_id}:agents"
            agents = self.client.smembers(agents_key)
            return list(agents) if agents else []
        except RedisError as e:
            print(f"Error listing agents from Redis: {e}")
            return []
    
    def read_shared_state(self, session_id: str) -> Optional[UnifiedState]:
        """Read shared state for a session."""
        try:
            key = f"uass:session:{session_id}:shared"
            value = self.client.get(key)
            
            if value is None:
                return None
            
            return UnifiedState.from_json(value)
        except (RedisError, json.JSONDecodeError) as e:
            print(f"Error reading shared state from Redis: {e}")
            return None
    
    def write_shared_state(self, session_id: str, state: UnifiedState) -> bool:
        """Write shared state for a session."""
        try:
            key = f"uass:session:{session_id}:shared"
            value = state.to_json()
            self.client.set(key, value)
            return True
        except RedisError as e:
            print(f"Error writing shared state to Redis: {e}")
            return False
    
    def get_history(self, session_id: str, agent_id: str, limit: int = 10) -> List[dict]:
        """Get state history for an agent."""
        try:
            history_key = f"uass:session:{session_id}:agent:{agent_id}:history"
            history = self.client.lrange(history_key, 0, limit - 1)
            return [json.loads(entry) for entry in history] if history else []
        except (RedisError, json.JSONDecodeError) as e:
            print(f"Error reading history from Redis: {e}")
            return []

