#!/usr/bin/env python3
"""
Conversation management for the agent
"""

from typing import Dict, List, Any
from datetime import datetime


class ConversationManager:
    """
    Manages conversation history and context
    """
    
    def __init__(self, max_history: int = 10):
        self.max_history = max_history
    
    def add_exchange(self, history: List[Dict[str, str]], user_query: str, agent_response: str) -> List[Dict[str, str]]:
        """
        Add a new user-agent exchange to the conversation history
        """
        updated_history = history.copy()
        
        # Add user message
        updated_history.append({
            "role": "user",
            "content": user_query,
            "timestamp": datetime.now().isoformat()
        })
        
        # Add agent response
        updated_history.append({
            "role": "assistant", 
            "content": agent_response,
            "timestamp": datetime.now().isoformat()
        })
        
        # Keep only the last max_history exchanges (pairs)
        if len(updated_history) > self.max_history * 2:
            updated_history = updated_history[-self.max_history * 2:]
        
        return updated_history
    
    def format_for_lightrag(self, history: List[Dict[str, str]]) -> List[Dict[str, str]]:
        """
        Format conversation history for LightRAG API
        """
        formatted = []
        for exchange in history:
            formatted.append({
                "role": exchange["role"],
                "content": exchange["content"]
            })
        return formatted
    
    def get_context_summary(self, history: List[Dict[str, str]]) -> str:
        """
        Generate a summary of recent conversation context
        """
        if not history:
            return "No previous conversation context."
        
        recent_exchanges = history[-6:]  # Last 3 exchanges
        context_parts = []
        
        for i, exchange in enumerate(recent_exchanges):
            role = "User" if exchange["role"] == "user" else "Assistant"
            context_parts.append(f"{role}: {exchange['content']}")
        
        return "\n".join(context_parts)


class ConversationFormatter:
    """
    Format conversation data for different purposes
    """
    
    @staticmethod
    def build_conversation_history(history: List[Dict[str, str]], max_turns: int = 10) -> List[Dict[str, str]]:
        """
        Build conversation history for LightRAG API
        """
        if not history:
            return []
        
        # Take last max_turns pairs (user + assistant)
        recent_history = history[-max_turns*2:]
        formatted = []
        
        for exchange in recent_history:
            # Handle both Message objects and dictionary formats
            if hasattr(exchange, 'role'):
                role = exchange.role
                content = exchange.content
            else:
                role = exchange["role"]
                content = exchange["content"]
            
            formatted.append({
                "role": role,
                "content": content
            })
        
        return formatted
    
    @staticmethod
    def create_context_summary(history: List[Dict[str, str]]) -> str:
        """
        Create a summary of conversation context
        """
        if not history:
            return "No previous conversation."
        
        recent_exchanges = history[-4:]  # Last 2 exchanges
        context_parts = []
        
        for exchange in recent_exchanges:
            # Handle both Message objects and dictionary formats
            if hasattr(exchange, 'role'):
                role = "User" if exchange.role == "user" else "Assistant"
                content = exchange.content
            else:
                role = "User" if exchange["role"] == "user" else "Assistant"
                content = exchange["content"]
            
            content = content[:100] + "..." if len(content) > 100 else content
            context_parts.append(f"{role}: {content}")
        
        return "\n".join(context_parts)