# src/utils/conversation_manager.py
from typing import List, Dict, Optional
import tiktoken
from datetime import datetime

class ConversationManager:
    def __init__(
        self,
        max_tokens: int = 4000,
        max_messages: int = 10,
        model: str = "gpt-3.5-turbo"
    ):
        """
        Initialize conversation manager
        
        Args:
            max_tokens (int): Maximum tokens to keep in context
            max_messages (int): Maximum number of messages to keep
            model (str): Model name for token counting
        """
        self.max_tokens = max_tokens
        self.max_messages = max_messages
        self.encoding = tiktoken.encoding_for_model(model)
        
    def format_messages(self, messages: List[Dict]) -> str:
        """Format messages into a conversation string"""
        formatted = []
        for msg in messages:
            role = msg.get('role', 'unknown')
            content = msg.get('content', '')
            formatted.append(f"{role.capitalize()}: {content}")
        return "\n".join(formatted)
    
    def count_tokens(self, text: str) -> int:
        """Count tokens in text"""
        return len(self.encoding.encode(text))
    
    def get_relevant_history(
        self,
        messages: List[Dict],
        current_query: str,
        max_tokens: Optional[int] = None
    ) -> List[Dict]:
        """
        Get relevant conversation history within token limit
        
        Args:
            messages (List[Dict]): Full message history
            current_query (str): Current user query
            max_tokens (Optional[int]): Override default max tokens
            
        Returns:
            List[Dict]: Relevant message history
        """
        max_tokens = max_tokens or self.max_tokens
        current_tokens = self.count_tokens(current_query)
        
        # Keep track of tokens and messages
        history = []
        total_tokens = current_tokens
        
        # Process messages from most recent to oldest
        for msg in reversed(messages[-self.max_messages:]):
            msg_text = f"{msg['role']}: {msg['content']}\n"
            msg_tokens = self.count_tokens(msg_text)
            
            # Check if adding this message would exceed token limit
            if total_tokens + msg_tokens > max_tokens:
                break
                
            total_tokens += msg_tokens
            history.append(msg)
        
        # Reverse back to chronological order
        return list(reversed(history))
    
    def generate_prompt_with_history(
        self,
        current_query: str,
        history: List[Dict],
        context_docs: List[str]
    ) -> str:
        """
        Generate a prompt that includes conversation history and context
        
        Args:
            current_query (str): Current user query
            history (List[Dict]): Relevant conversation history
            context_docs (List[str]): Retrieved context documents
            
        Returns:
            str: Formatted prompt
        """
        # Format conversation history
        conversation_context = self.format_messages(history)
        
        # Format context documents
        context_str = "\n\n".join(context_docs)
        
        prompt = f"""
Previous Conversation:
{conversation_context}

Relevant Context:
{context_str}

Current Query: {current_query}

Based on the previous conversation and the provided context, please provide a comprehensive and accurate response that maintains continuity with the conversation history."""
        
        return prompt