"""
Conversation Management
Handles message history and context window management.
"""
from typing import List, Dict, Optional
from datetime import datetime
import uuid
from loguru import logger


class ConversationManager:
    """
    Manages conversation history and context.
    
    Features:
    - Track multiple conversations
    - Automatic context window management
    - Message summarization when context grows too large
    """
    
    def __init__(self, max_messages: int = 50, max_tokens: int = 8000):
        self.conversations: Dict[str, List[Dict]] = {}
        self.max_messages = max_messages
        self.max_tokens = max_tokens
    
    def create_conversation(self) -> str:
        """Create a new conversation and return its ID."""
        conv_id = str(uuid.uuid4())
        self.conversations[conv_id] = []
        logger.debug(f"Created conversation: {conv_id}")
        return conv_id
    
    def get_conversation(self, conv_id: str) -> List[Dict]:
        """Get messages for a conversation."""
        return self.conversations.get(conv_id, [])
    
    def add_message(
        self,
        conv_id: str,
        role: str,
        content: str,
        metadata: Optional[Dict] = None
    ) -> None:
        """Add a message to a conversation."""
        if conv_id not in self.conversations:
            self.conversations[conv_id] = []
        
        message = {
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat(),
        }
        
        if metadata:
            message["metadata"] = metadata
        
        self.conversations[conv_id].append(message)
        
        # Trim if needed
        self._trim_conversation(conv_id)
    
    def _trim_conversation(self, conv_id: str) -> None:
        """Trim conversation if it exceeds limits."""
        messages = self.conversations.get(conv_id, [])
        
        if len(messages) > self.max_messages:
            # Keep system messages and last N messages
            system_messages = [m for m in messages if m["role"] == "system"]
            other_messages = [m for m in messages if m["role"] != "system"]
            
            # Keep last N-1 messages (plus system)
            keep_count = self.max_messages - len(system_messages) - 1
            trimmed = system_messages + other_messages[-keep_count:]
            
            self.conversations[conv_id] = trimmed
            logger.debug(f"Trimmed conversation {conv_id} to {len(trimmed)} messages")
    
    def delete_conversation(self, conv_id: str) -> None:
        """Delete a conversation."""
        if conv_id in self.conversations:
            del self.conversations[conv_id]
            logger.debug(f"Deleted conversation: {conv_id}")
    
    def list_conversations(self) -> List[str]:
        """List all conversation IDs."""
        return list(self.conversations.keys())
    
    def estimate_tokens(self, messages: List[Dict]) -> int:
        """Estimate token count for messages."""
        # Rough estimate: ~4 characters per token
        total_chars = sum(
            len(m.get("content", "")) + len(m.get("role", ""))
            for m in messages
        )
        return total_chars // 4
