"""Storage backends for conversation persistence."""

import asyncio
import json
import os
import pickle
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Dict, List, Optional

from justllms.conversations.models import ConversationMessage, ConversationSummary, ConversationAnalytics


class ConversationStorage(ABC):
    """Abstract base class for conversation storage backends."""
    
    @abstractmethod
    async def save_conversation(
        self, 
        conversation_id: str, 
        messages: List[ConversationMessage],
        summary: ConversationSummary,
        analytics: ConversationAnalytics,
        metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        """Save a complete conversation."""
        pass
    
    @abstractmethod
    async def load_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]:
        """Load a conversation by ID."""
        pass
    
    @abstractmethod
    async def list_conversations(
        self, 
        limit: Optional[int] = None,
        offset: int = 0,
        filters: Optional[Dict[str, Any]] = None
    ) -> List[ConversationSummary]:
        """List conversations with optional filtering."""
        pass
    
    @abstractmethod
    async def delete_conversation(self, conversation_id: str) -> bool:
        """Delete a conversation."""
        pass
    
    @abstractmethod
    async def save_message(self, conversation_id: str, message: ConversationMessage) -> None:
        """Save a single message to a conversation."""
        pass
    
    @abstractmethod
    async def get_messages(
        self, 
        conversation_id: str,
        limit: Optional[int] = None,
        offset: int = 0
    ) -> List[ConversationMessage]:
        """Get messages from a conversation."""
        pass


class MemoryStorage(ConversationStorage):
    """In-memory storage for conversations."""
    
    def __init__(self):
        self.conversations: Dict[str, Dict[str, Any]] = {}
        self._lock = asyncio.Lock()
    
    async def save_conversation(
        self,
        conversation_id: str,
        messages: List[ConversationMessage],
        summary: ConversationSummary,
        analytics: ConversationAnalytics,
        metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        """Save a complete conversation."""
        async with self._lock:
            self.conversations[conversation_id] = {
                "messages": [msg.model_dump() for msg in messages],
                "summary": summary.model_dump(),
                "analytics": analytics.model_dump(),
                "metadata": metadata or {}
            }
    
    async def load_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]:
        """Load a conversation by ID."""
        async with self._lock:
            data = self.conversations.get(conversation_id)
            if not data:
                return None
            
            # Reconstruct objects
            messages = [ConversationMessage(**msg) for msg in data["messages"]]
            summary = ConversationSummary(**data["summary"])
            analytics = ConversationAnalytics(**data["analytics"])
            
            return {
                "messages": messages,
                "summary": summary,
                "analytics": analytics,
                "metadata": data.get("metadata", {})
            }
    
    async def list_conversations(
        self,
        limit: Optional[int] = None,
        offset: int = 0,
        filters: Optional[Dict[str, Any]] = None
    ) -> List[ConversationSummary]:
        """List conversations with optional filtering."""
        async with self._lock:
            summaries = []
            for conv_data in self.conversations.values():
                summary = ConversationSummary(**conv_data["summary"])
                
                # Apply filters if provided
                if filters:
                    if filters.get("state") and summary.state != filters["state"]:
                        continue
                    if filters.get("tags") and not any(tag in summary.tags for tag in filters["tags"]):
                        continue
                
                summaries.append(summary)
            
            # Sort by updated_at descending
            summaries.sort(key=lambda x: x.updated_at, reverse=True)
            
            # Apply pagination
            if limit:
                summaries = summaries[offset:offset + limit]
            
            return summaries
    
    async def delete_conversation(self, conversation_id: str) -> bool:
        """Delete a conversation."""
        async with self._lock:
            if conversation_id in self.conversations:
                del self.conversations[conversation_id]
                return True
            return False
    
    async def save_message(self, conversation_id: str, message: ConversationMessage) -> None:
        """Save a single message to a conversation."""
        async with self._lock:
            if conversation_id in self.conversations:
                self.conversations[conversation_id]["messages"].append(message.model_dump())
    
    async def get_messages(
        self,
        conversation_id: str,
        limit: Optional[int] = None,
        offset: int = 0
    ) -> List[ConversationMessage]:
        """Get messages from a conversation."""
        async with self._lock:
            if conversation_id not in self.conversations:
                return []
            
            messages_data = self.conversations[conversation_id]["messages"]
            messages = [ConversationMessage(**msg) for msg in messages_data]
            
            if limit:
                messages = messages[offset:offset + limit]
            
            return messages


class DiskStorage(ConversationStorage):
    """Disk-based storage for conversations."""
    
    def __init__(self, storage_dir: Optional[Path] = None):
        self.storage_dir = storage_dir or Path.home() / ".justllms" / "conversations"
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self._lock = asyncio.Lock()
    
    def _get_conversation_path(self, conversation_id: str) -> Path:
        """Get file path for a conversation."""
        return self.storage_dir / f"{conversation_id}.json"
    
    async def save_conversation(
        self,
        conversation_id: str,
        messages: List[ConversationMessage],
        summary: ConversationSummary,
        analytics: ConversationAnalytics,
        metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        """Save a complete conversation."""
        data = {
            "messages": [msg.model_dump() for msg in messages],
            "summary": summary.model_dump(),
            "analytics": analytics.model_dump(),
            "metadata": metadata or {}
        }
        
        async with self._lock:
            file_path = self._get_conversation_path(conversation_id)
            await asyncio.get_event_loop().run_in_executor(
                None, 
                lambda: file_path.write_text(json.dumps(data, indent=2))
            )
    
    async def load_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]:
        """Load a conversation by ID."""
        file_path = self._get_conversation_path(conversation_id)
        
        if not file_path.exists():
            return None
        
        async with self._lock:
            try:
                data = await asyncio.get_event_loop().run_in_executor(
                    None,
                    lambda: json.loads(file_path.read_text())
                )
                
                # Reconstruct objects
                messages = [ConversationMessage(**msg) for msg in data["messages"]]
                summary = ConversationSummary(**data["summary"])
                analytics = ConversationAnalytics(**data["analytics"])
                
                return {
                    "messages": messages,
                    "summary": summary,
                    "analytics": analytics,
                    "metadata": data.get("metadata", {})
                }
            except (json.JSONDecodeError, KeyError) as e:
                # Error loading conversation - may not exist
                return None
    
    async def list_conversations(
        self,
        limit: Optional[int] = None,
        offset: int = 0,
        filters: Optional[Dict[str, Any]] = None
    ) -> List[ConversationSummary]:
        """List conversations with optional filtering."""
        summaries = []
        
        async with self._lock:
            for file_path in self.storage_dir.glob("*.json"):
                try:
                    data = await asyncio.get_event_loop().run_in_executor(
                        None,
                        lambda fp=file_path: json.loads(fp.read_text())
                    )
                    
                    summary = ConversationSummary(**data["summary"])
                    
                    # Apply filters if provided
                    if filters:
                        if filters.get("state") and summary.state != filters["state"]:
                            continue
                        if filters.get("tags") and not any(tag in summary.tags for tag in filters["tags"]):
                            continue
                    
                    summaries.append(summary)
                    
                except (json.JSONDecodeError, KeyError):
                    continue
        
        # Sort by updated_at descending
        summaries.sort(key=lambda x: x.updated_at, reverse=True)
        
        # Apply pagination
        if limit:
            summaries = summaries[offset:offset + limit]
        
        return summaries
    
    async def delete_conversation(self, conversation_id: str) -> bool:
        """Delete a conversation."""
        file_path = self._get_conversation_path(conversation_id)
        
        async with self._lock:
            try:
                await asyncio.get_event_loop().run_in_executor(None, file_path.unlink)
                return True
            except FileNotFoundError:
                return False
    
    async def save_message(self, conversation_id: str, message: ConversationMessage) -> None:
        """Save a single message to a conversation."""
        # For disk storage, we need to load, modify, and save the entire conversation
        conversation_data = await self.load_conversation(conversation_id)
        if conversation_data:
            conversation_data["messages"].append(message)
            await self.save_conversation(
                conversation_id,
                conversation_data["messages"],
                conversation_data["summary"], 
                conversation_data["analytics"],
                conversation_data["metadata"]
            )
    
    async def get_messages(
        self,
        conversation_id: str,
        limit: Optional[int] = None,
        offset: int = 0
    ) -> List[ConversationMessage]:
        """Get messages from a conversation."""
        conversation_data = await self.load_conversation(conversation_id)
        if not conversation_data:
            return []
        
        messages = conversation_data["messages"]
        
        if limit:
            messages = messages[offset:offset + limit]
        
        return messages


class RedisStorage(ConversationStorage):
    """Redis-based storage for conversations."""
    
    def __init__(
        self,
        host: str = "localhost",
        port: int = 6379,
        db: int = 0,
        password: Optional[str] = None,
        prefix: str = "justllms:conversations:",
        **redis_kwargs
    ):
        self.host = host
        self.port = port
        self.db = db
        self.password = password
        self.prefix = prefix
        self.redis_kwargs = redis_kwargs
        
        try:
            import redis.asyncio as aioredis
            self.redis_available = True
            
            self.redis = aioredis.Redis(
                host=host,
                port=port,
                db=db,
                password=password,
                **redis_kwargs
            )
        except ImportError:
            self.redis_available = False
            raise ImportError("Redis support requires the 'redis' package. Install with: pip install redis")
    
    def _make_key(self, key: str) -> str:
        """Add prefix to key."""
        return f"{self.prefix}{key}"
    
    def _serialize(self, data: Any) -> bytes:
        """Serialize data for Redis storage."""
        return pickle.dumps(data)
    
    def _deserialize(self, data: bytes) -> Any:
        """Deserialize data from Redis storage."""
        return pickle.loads(data)
    
    async def save_conversation(
        self,
        conversation_id: str,
        messages: List[ConversationMessage],
        summary: ConversationSummary,
        analytics: ConversationAnalytics,
        metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        """Save a complete conversation."""
        data = {
            "messages": [msg.model_dump() for msg in messages],
            "summary": summary.model_dump(),
            "analytics": analytics.model_dump(),
            "metadata": metadata or {}
        }
        
        key = self._make_key(conversation_id)
        serialized_data = self._serialize(data)
        
        await self.redis.set(key, serialized_data)
    
    async def load_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]:
        """Load a conversation by ID."""
        key = self._make_key(conversation_id)
        data = await self.redis.get(key)
        
        if not data:
            return None
        
        try:
            deserialized_data = self._deserialize(data)
            
            # Reconstruct objects
            messages = [ConversationMessage(**msg) for msg in deserialized_data["messages"]]
            summary = ConversationSummary(**deserialized_data["summary"])
            analytics = ConversationAnalytics(**deserialized_data["analytics"])
            
            return {
                "messages": messages,
                "summary": summary,
                "analytics": analytics,
                "metadata": deserialized_data.get("metadata", {})
            }
        except Exception as e:
            # Error deserializing conversation - corrupted data
            return None
    
    async def list_conversations(
        self,
        limit: Optional[int] = None,
        offset: int = 0,
        filters: Optional[Dict[str, Any]] = None
    ) -> List[ConversationSummary]:
        """List conversations with optional filtering."""
        pattern = f"{self.prefix}*"
        keys = []
        
        async for key in self.redis.scan_iter(match=pattern):
            keys.append(key.decode('utf-8'))
        
        summaries = []
        
        for key in keys:
            data = await self.redis.get(key)
            if data:
                try:
                    deserialized_data = self._deserialize(data)
                    summary = ConversationSummary(**deserialized_data["summary"])
                    
                    # Apply filters if provided
                    if filters:
                        if filters.get("state") and summary.state != filters["state"]:
                            continue
                        if filters.get("tags") and not any(tag in summary.tags for tag in filters["tags"]):
                            continue
                    
                    summaries.append(summary)
                except Exception:
                    continue
        
        # Sort by updated_at descending
        summaries.sort(key=lambda x: x.updated_at, reverse=True)
        
        # Apply pagination
        if limit:
            summaries = summaries[offset:offset + limit]
        
        return summaries
    
    async def delete_conversation(self, conversation_id: str) -> bool:
        """Delete a conversation."""
        key = self._make_key(conversation_id)
        result = await self.redis.delete(key)
        return result > 0
    
    async def save_message(self, conversation_id: str, message: ConversationMessage) -> None:
        """Save a single message to a conversation."""
        # For Redis, we need to load, modify, and save the entire conversation
        conversation_data = await self.load_conversation(conversation_id)
        if conversation_data:
            conversation_data["messages"].append(message)
            await self.save_conversation(
                conversation_id,
                conversation_data["messages"],
                conversation_data["summary"],
                conversation_data["analytics"],
                conversation_data["metadata"]
            )
    
    async def get_messages(
        self,
        conversation_id: str,
        limit: Optional[int] = None,
        offset: int = 0
    ) -> List[ConversationMessage]:
        """Get messages from a conversation."""
        conversation_data = await self.load_conversation(conversation_id)
        if not conversation_data:
            return []
        
        messages = conversation_data["messages"]
        
        if limit:
            messages = messages[offset:offset + limit]
        
        return messages