cache/conversation.py

from logger.logger import get_logger
logger = get_logger(__name__)
from typing import List, Dict, Any, Optional
from datetime import datetime, timedelta
import asyncio
import uuid
import json
from services.conversation_service import conversation_db
from cache.redis_client import redis_client

class ConversationCache:
    def __init__(self, ttl_hours: int = 2, batch_size: int = 50):
        self.ttl_seconds = ttl_hours * 3600
        self.delete_after_seconds = 3600
        self.batch_size = batch_size
        self._lock = asyncio.Lock()
        self._cleanup_task = None
        self._shutdown_event = asyncio.Event()
        self._save_queue: asyncio.Queue = asyncio.Queue()
        self._worker_task = None
        self._is_running = True
        self._redis = redis_client
        self._redis_initialized = False
    
    def _ensure_redis(self):
        if not self._redis_initialized:
            self._redis.connect()
            self._redis_initialized = True
    
    def _get_conv_key(self, conv_id: str) -> str:
        return f"conv:{conv_id}"
    
    def _get_user_key(self, user_id: str) -> str:
        return f"user:{user_id}:convs"
    
    def _get_msgs_key(self, conv_id: str) -> str:
        return f"conv:{conv_id}:msgs"
    
    def _get_saved_key(self, conv_id: str) -> str:
        return f"conv:{conv_id}:saved"
    
    def _get_status_key(self, conv_id: str) -> str:
        return f"conv:{conv_id}:status"
    
    def _get_msg_id(self, message: Dict[str, Any]) -> str:
        ts = message['timestamp'].isoformat() if isinstance(message['timestamp'], datetime) else message['timestamp']
        return f"{message['role']}:{message['content'][:50]}:{ts}"
    
    def _format_response(self, conv: Dict[str, Any], include_messages: bool, max_messages: Optional[int]) -> Dict[str, Any]:
        conv_copy = conv.copy()
        for field in ["expires_at", "created_at", "updated_at"]:
            if field in conv_copy and isinstance(conv_copy[field], datetime):
                conv_copy[field] = conv_copy[field].isoformat()
        
        if not include_messages:
            conv_copy.pop("messages", None)
        elif max_messages and conv_copy.get("messages"):
            conv_copy["messages"] = conv_copy["messages"][-max_messages:]
            for msg in conv_copy["messages"]:
                if isinstance(msg["timestamp"], datetime):
                    msg["timestamp"] = msg["timestamp"].isoformat()
        return conv_copy
    
    async def _load_from_redis(self, conv_id: str) -> Optional[Dict[str, Any]]:
        try:
            self._ensure_redis()
            conv_data = self._redis.hgetall(self._get_conv_key(conv_id))
            if not conv_data:
                return None
            
            raw_msgs = self._redis.zrange(self._get_msgs_key(conv_id), 0, -1)
            messages = []
            for raw in raw_msgs:
                msg = json.loads(raw)
                msg["timestamp"] = datetime.fromisoformat(msg["timestamp"])
                msg["metadata"] = json.loads(msg["metadata"])
                messages.append(msg)
            
            return {
                "conversation_id": conv_data["conversation_id"],
                "user_id": conv_data["user_id"],
                "metadata": json.loads(conv_data["metadata"]),
                "created_at": datetime.fromisoformat(conv_data["created_at"]),
                "updated_at": datetime.fromisoformat(conv_data["updated_at"]),
                "expires_at": datetime.fromisoformat(conv_data["expires_at"]),
                "messages": messages
            }
        except Exception as e:
            logger.error(f"Error loading from Redis: {e}")
            return None
    
    async def _save_to_redis(self, conv_data: Dict[str, Any]) -> bool:
        try:
            self._ensure_redis()
            now = datetime.now()
            pipe = self._redis.pipeline()
            
            pipe.hset(self._get_conv_key(conv_data["conversation_id"]), mapping={
                "conversation_id": conv_data["conversation_id"],
                "user_id": conv_data["user_id"],
                "metadata": json.dumps(conv_data["metadata"]),
                "created_at": conv_data["created_at"].isoformat(),
                "updated_at": conv_data["updated_at"].isoformat(),
                "expires_at": (now + timedelta(seconds=self.ttl_seconds)).isoformat()
            })
            pipe.expire(self._get_conv_key(conv_data["conversation_id"]), self.ttl_seconds)
            
            pipe.lpush(self._get_user_key(conv_data["user_id"]), conv_data["conversation_id"])
            pipe.expire(self._get_user_key(conv_data["user_id"]), self.ttl_seconds)
            
            for msg in conv_data.get("messages", []):
                msg_data = {
                    "role": msg["role"],
                    "content": msg["content"],
                    "metadata": json.dumps(msg.get("metadata", {})),
                    "timestamp": msg["timestamp"].isoformat() if isinstance(msg["timestamp"], datetime) else msg["timestamp"]
                }
                ts = msg["timestamp"].timestamp() if isinstance(msg["timestamp"], datetime) else datetime.fromisoformat(msg["timestamp"]).timestamp()
                pipe.zadd(self._get_msgs_key(conv_data["conversation_id"]), {json.dumps(msg_data): ts})
            
            pipe.expire(self._get_msgs_key(conv_data["conversation_id"]), self.ttl_seconds)
            pipe.set(self._get_status_key(conv_data["conversation_id"]), "restored", ex=self.ttl_seconds)
            pipe.execute()
            return True
        except Exception as e:
            logger.error(f"Error saving to Redis: {e}")
            return False
    
    async def _db_worker(self):
        while self._is_running or not self._save_queue.empty():
            try:
                # Wait until an item is available
                conv_id, data = await self._save_queue.get()

                try:
                    success = False

                    if (
                        isinstance(data, dict)
                        and data.get("action") == "create_conversation"
                    ):
                        conv_data = await self._load_from_redis(conv_id)

                        if conv_data:
                            success = await conversation_db.save_conversation(conv_data)

                    elif (
                        isinstance(data, dict)
                        and data.get("action") == "update_metadata"
                    ):
                        conv_data = await self._load_from_redis(conv_id)

                        if conv_data:
                            success = await conversation_db.save_conversation(conv_data)

                    else:
                        success = await conversation_db.save_message(conv_id, data)

                    if success:
                        logger.debug(f"Successfully saved item for conversation {conv_id}")
                    else:
                        logger.warning(f"Failed to save item for conversation {conv_id}")

                except Exception as e:
                    logger.exception(f"Error processing conversation {conv_id}: {e}")

                finally:
                    # Mark queue item as completed ONLY after DB work finishes
                    self._save_queue.task_done()

            except asyncio.CancelledError:
                logger.info("DB worker cancelled")
                break

            except Exception as e:
                logger.exception(f"Unexpected error in DB worker: {e}")
                await asyncio.sleep(1)
    
    async def create_conversation(self, user_id: str, conversation_id: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None) -> str:
        async with self._lock:
            if conversation_id is None:
                conversation_id = str(uuid.uuid4())
            
            now = datetime.now()
            self._ensure_redis()
            
            # Save to Redis immediately
            pipe = self._redis.pipeline()
            pipe.hset(self._get_conv_key(conversation_id), mapping={
                "conversation_id": conversation_id,
                "user_id": user_id,
                "metadata": json.dumps(metadata or {}),
                "created_at": now.isoformat(),
                "updated_at": now.isoformat(),
                "expires_at": (now + timedelta(seconds=self.ttl_seconds)).isoformat()
            })
            pipe.expire(self._get_conv_key(conversation_id), self.ttl_seconds)
            pipe.lpush(self._get_user_key(user_id), conversation_id)
            pipe.expire(self._get_user_key(user_id), self.ttl_seconds)
            pipe.set(self._get_status_key(conversation_id), "pending", ex=self.ttl_seconds)
            pipe.execute()
            
            # Queue for background DB save
            await self._save_queue.put((conversation_id, {"action": "create_conversation"}))
            
            logger.info(f"Created conversation {conversation_id} for user {user_id} in Redis")
            return conversation_id
    
    async def add_message(self, conversation_id: str, role: str, content: str, metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        async with self._lock:
            # Check Redis first
            self._ensure_redis()
            if not self._redis.exists(self._get_conv_key(conversation_id)):
                # If not in Redis, try to load from DB and restore to Redis
                conv_data = await conversation_db.load_conversation(conversation_id)
                if not conv_data:
                    raise ValueError(f"Conversation {conversation_id} not found")
                await self._save_to_redis(conv_data)
            
            now = datetime.now()
            message = {"role": role, "content": content, "timestamp": now, "metadata": metadata or {}}
            
            # Save to Redis immediately
            pipe = self._redis.pipeline()
            pipe.expire(self._get_conv_key(conversation_id), self.ttl_seconds)
            msg_data = {"role": role, "content": content, "metadata": json.dumps(metadata or {}), "timestamp": now.isoformat()}
            pipe.zadd(self._get_msgs_key(conversation_id), {json.dumps(msg_data): now.timestamp()})
            pipe.expire(self._get_msgs_key(conversation_id), self.ttl_seconds)
            pipe.hset(self._get_conv_key(conversation_id), "updated_at", now.isoformat())
            pipe.hset(self._get_conv_key(conversation_id), "expires_at", (now + timedelta(seconds=self.ttl_seconds)).isoformat())
            pipe.execute()
            
            # Queue for background DB save
            await self._save_queue.put((conversation_id, message))
            
            logger.debug(f"Added {role} message to conversation {conversation_id} in Redis")
            return message
    
    async def get_conversation(self, conversation_id: str, include_messages: bool = True, max_messages: Optional[int] = None) -> Optional[Dict[str, Any]]:
        async with self._lock:
            # Always read from Redis first
            conv_data = await self._load_from_redis(conversation_id)
            if conv_data:
                logger.debug(f"Conversation {conversation_id} found in Redis")
                return self._format_response(conv_data, include_messages, max_messages)
            
            # If not in Redis, check Database
            logger.debug(f"Conversation {conversation_id} not in Redis, checking DB")
            conv_data = await conversation_db.load_conversation(conversation_id)
            if conv_data:
                # Restore to Redis for future requests
                await self._save_to_redis(conv_data)
                return self._format_response(conv_data, include_messages, max_messages)
            
            return None
    
    async def get_user_conversations(self, user_id: str, limit: int = 10, include_messages: bool = False) -> List[Dict[str, Any]]:
        async with self._lock:
            # Always read from Redis first
            self._ensure_redis()
            conv_ids = self._redis.lrange(self._get_user_key(user_id), 0, limit - 1)
            
            if conv_ids:
                logger.debug(f"User {user_id} conversations found in Redis")
                convs = []
                for conv_id in conv_ids:
                    conv = await self.get_conversation(conv_id, include_messages)
                    if conv:
                        convs.append(conv)
                return convs
            
            # If not in Redis, check Database
            logger.debug(f"User {user_id} conversations not in Redis, checking DB")
            convs = await conversation_db.load_user_conversations(user_id, limit)
            # Restore to Redis
            for conv in convs:
                await self._save_to_redis(conv)
            return convs
    
    async def update_conversation_metadata(self, conversation_id: str, metadata: Dict[str, Any]) -> bool:
        async with self._lock:
            # Check Redis first
            conv_data = await self.get_conversation(conversation_id, include_messages=True)
            if not conv_data:
                return False
            
            conv_data["metadata"].update(metadata)
            conv_data["updated_at"] = datetime.now()
            
            # Update in Redis immediately
            self._ensure_redis()
            pipe = self._redis.pipeline()
            pipe.hset(self._get_conv_key(conversation_id), "metadata", json.dumps(conv_data["metadata"]))
            pipe.hset(self._get_conv_key(conversation_id), "updated_at", datetime.now().isoformat())
            pipe.execute()
            
            # Queue for background DB save
            await self._save_queue.put((conversation_id, {"action": "update_metadata", "metadata": metadata}))
            
            logger.debug(f"Updated metadata for conversation {conversation_id}")
            return True
    
    async def delete_conversation(self, conversation_id: str) -> bool:
        async with self._lock:
            # Check Redis first
            conv_data = await self.get_conversation(conversation_id, include_messages=True)
            if not conv_data:
                return False
            
            # Flush pending saves to DB
            await self.flush_all_to_db()
            
            # Save to DB if not already saved
            status = self._redis.get(self._get_status_key(conversation_id))
            if status != "saved":
                await conversation_db.save_conversation(conv_data)
            
            # Delete from Redis
            self._ensure_redis()
            pipe = self._redis.pipeline()
            pipe.delete(self._get_conv_key(conversation_id))
            pipe.delete(self._get_msgs_key(conversation_id))
            pipe.delete(self._get_saved_key(conversation_id))
            pipe.delete(self._get_status_key(conversation_id))
            pipe.lrem(self._get_user_key(conv_data["user_id"]), 0, conversation_id)
            pipe.execute()
            
            logger.info(f"Deleted conversation {conversation_id} from Redis")
            return True
    
    async def delete_user_conversations(self, user_id: str) -> int:
        async with self._lock:
            self._ensure_redis()
            conv_ids = self._redis.lrange(self._get_user_key(user_id), 0, -1)
            count = 0
            
            for conv_id in conv_ids:
                conv_data = await self.get_conversation(conv_id, include_messages=True)
                if conv_data:
                    status = self._redis.get(self._get_status_key(conv_id))
                    if status != "saved":
                        await conversation_db.save_conversation(conv_data)
                    
                    pipe = self._redis.pipeline()
                    pipe.delete(self._get_conv_key(conv_id))
                    pipe.delete(self._get_msgs_key(conv_id))
                    pipe.delete(self._get_saved_key(conv_id))
                    pipe.delete(self._get_status_key(conv_id))
                    pipe.execute()
                    count += 1
            
            self._redis.delete(self._get_user_key(user_id))
            await self.flush_all_to_db()
            logger.info(f"Deleted {count} conversations for user {user_id}")
            return count
    
    async def get_conversation_history(self, conversation_id: str, max_messages: int = 20, format_type: str = "list") -> List:
        conv = await self.get_conversation(conversation_id, include_messages=True, max_messages=max_messages)
        if not conv:
            return []
        
        messages_list = conv.get("messages", [])
        if format_type == "string":
            return [f"{msg['role']}: {msg['content']}" for msg in messages_list]
        return [{"role": msg["role"], "content": msg["content"]} for msg in messages_list]
    
    async def clear_expired(self) -> int:
        self._ensure_redis()
        pattern = "conv:*"
        cursor = 0
        deleted = 0
        
        while True:
            cursor, keys = self._redis.get_client().scan(cursor, match=pattern, count=100)
            for key in keys:
                if ":" in key:
                    conv_id = key.split(":")[1]
                    created_at_str = self._redis.hget(key, "created_at")
                    if created_at_str:
                        created_at = datetime.fromisoformat(created_at_str)
                        age = (datetime.now() - created_at).total_seconds()
                        
                        if age > self.delete_after_seconds:
                            status = self._redis.get(self._get_status_key(conv_id))
                            if status == "saved":
                                user_id = self._redis.hget(key, "user_id")
                                pipe = self._redis.pipeline()
                                pipe.delete(key)
                                pipe.delete(self._get_msgs_key(conv_id))
                                pipe.delete(self._get_saved_key(conv_id))
                                pipe.delete(self._get_status_key(conv_id))
                                if user_id:
                                    pipe.lrem(self._get_user_key(user_id), 0, conv_id)
                                pipe.execute()
                                deleted += 1
            if cursor == 0:
                break
        
        if deleted > 0:
            logger.info(f"Cleaned up {deleted} expired conversations from Redis")
        return deleted
    
    async def get_stats(self) -> Dict[str, Any]:
        self._ensure_redis()
        pattern = "conv:*"
        cursor = 0
        conv_count = 0
        msg_count = 0
        
        while True:
            cursor, keys = self._redis.get_client().scan(cursor, match=pattern, count=100)
            for key in keys:
                if ":" in key:
                    conv_id = key.split(":")[1]
                    if conv_id:
                        conv_count += 1
                        msg_count += self._redis.zcard(self._get_msgs_key(conv_id))
            if cursor == 0:
                break
        
        return {
            "total_conversations": conv_count,
            "total_messages": msg_count,
            "ttl_hours": self.ttl_seconds // 3600,
            "delete_after_hours": self.delete_after_seconds // 3600,
            "queue_size": self._save_queue.qsize(),
            "redis_connected": self._redis_initialized and self._redis.ping()
        }
    
    async def start_background_tasks(self):
        await self.start_cleanup_task()
        if self._worker_task is None or self._worker_task.done():
            self._is_running = True
            self._worker_task = asyncio.create_task(self._db_worker())
            logger.info("Started DB worker task")
    
    async def stop_background_tasks(self):
        await self.stop_cleanup_task()
        self._is_running = False
        if self._worker_task and not self._worker_task.done():
            self._worker_task.cancel()
            try:
                await self._worker_task
            except asyncio.CancelledError:
                pass
            logger.info("Stopped DB worker task")
    
    async def start_cleanup_task(self):
        if self._cleanup_task is None or self._cleanup_task.done():
            self._shutdown_event.clear()
            self._cleanup_task = asyncio.create_task(self._cleanup_loop())
            logger.info("Started background cleanup task")
    
    async def stop_cleanup_task(self):
        if self._cleanup_task and not self._cleanup_task.done():
            self._shutdown_event.set()
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
            logger.info("Stopped background cleanup task")
    
    async def _cleanup_loop(self):
        try:
            while not self._shutdown_event.is_set():
                try:
                    await asyncio.wait_for(self._shutdown_event.wait(), timeout=3600)
                    break
                except asyncio.TimeoutError:
                    removed = await self.clear_expired()
                    if removed > 0:
                        logger.info(f"Cleanup job removed {removed} expired conversations")
        except asyncio.CancelledError:
            logger.info("Cleanup task received cancellation signal")
            raise
        except Exception as e:
            logger.error(f"Error in cleanup loop: {e}", exc_info=True)
    
    async def flush_all_to_db(self) -> int:
        remaining = self._save_queue.qsize()
        if remaining > 0:
            logger.info(f"Flushing {remaining} items to database...")
            await self._save_queue.join()
            logger.info(f"Flushed all {remaining} items to database")
        return remaining
    
    async def save_all_conversations_to_db(self) -> int:
        await self.flush_all_to_db()
        self._ensure_redis()
        pattern = "conv:*"
        cursor = 0
        saved = 0
        
        while True:
            cursor, keys = self._redis.get_client().scan(cursor, match=pattern, count=100)
            for key in keys:
                if ":" in key:
                    conv_id = key.split(":")[1]
                    if conv_id:
                        conv_data = await self._load_from_redis(conv_id)
                        if conv_data:
                            if await conversation_db.save_conversation(conv_data):
                                saved += 1
            if cursor == 0:
                break
        
        logger.info(f"Saved {saved} conversations to DB during shutdown")
        return saved
    
    async def close(self):
        self._redis.close()

cache = ConversationCache(ttl_hours=2)

async def initialize_cache():
    await cache.start_background_tasks()
    cache._ensure_redis()