#!/usr/bin/env python3
"""
Agent Memory Search - Optimized for Chat Agents

This class is designed to be initialized once and used repeatedly
to get similar memories for user queries in chat agents.
"""

import asyncio
import logging
from typing import List, Dict, Any, Optional
from memory_search import MemorySearcher


class AgentMemorySearch:
    """
    Memory search class optimized for agentic chat systems.
    
    Initialize once, then call get_similar_memories() for each user query.
    """
    
    def __init__(self, 
                 default_max_results: int = 10,
                 quality_threshold: float = 0.5,
                 include_conversations: bool = True,
                 include_entities: bool = True):
        """
        Initialize the agent memory search.
        
        Args:
            default_max_results: Default number of results to return
            quality_threshold: Minimum relevance score threshold (0.0-1.0)
            include_conversations: Whether to include conversation memories
            include_entities: Whether to include entity information
        """
        self.searcher = MemorySearcher()
        self.default_max_results = default_max_results
        self.quality_threshold = quality_threshold
        self.include_conversations = include_conversations
        self.include_entities = include_entities
        self.is_initialized = False
        
        # Setup logging
        self.logger = logging.getLogger(__name__)
    
    async def initialize(self):
        """Initialize the memory search client. Call this once at startup."""
        if not self.is_initialized:
            try:
                await self.searcher._ensure_client_initialized()
                self.is_initialized = True
                self.logger.info("Agent memory search initialized successfully")
            except Exception as e:
                self.logger.error(f"Failed to initialize memory search: {e}")
                raise
    
    def _calculate_relevance_score(self, rank: int, total_results: int) -> float:
        """Calculate relevance score based on ranking position."""
        if total_results == 0:
            return 0.0
        return 1.0 - (0.9 * (rank - 1) / max(1, total_results - 1))
    
    def _extract_fact_text(self, fact_dict: Dict) -> str:
        """Extract fact text from Graphiti edge/fact structure."""
        # Try different possible fields for the fact text
        if 'fact' in fact_dict:
            return fact_dict['fact']
        elif 'content' in fact_dict:
            return fact_dict['content']
        elif 'text' in fact_dict:
            return fact_dict['text']
        elif 'description' in fact_dict:
            return fact_dict['description']
        else:
            # Fallback: try to construct from source and target
            source = fact_dict.get('source_node_name', 'Unknown')
            target = fact_dict.get('target_node_name', 'Unknown')
            relation = fact_dict.get('relation_type', 'relates to')
            return f"{source} {relation} {target}"
    
    def _extract_entity_info(self, entity_dict: Dict) -> Dict[str, str]:
        """Extract entity name and summary from Graphiti node structure."""
        name = entity_dict.get('name', entity_dict.get('node_name', 'Unknown Entity'))
        
        # Try to get summary from different possible fields
        summary = (entity_dict.get('summary') or 
                  entity_dict.get('description') or 
                  entity_dict.get('content') or 
                  entity_dict.get('attributes', {}).get('summary') or
                  entity_dict.get('attributes', {}).get('description') or
                  'No description available')
        
        return {'name': name, 'summary': summary}
    
    def _extract_conversation_content(self, conversation_dict: Dict) -> str:
        """Extract conversation content from Graphiti episode structure."""
        # Try different possible fields for conversation content
        if 'content' in conversation_dict:
            return conversation_dict['content']
        elif 'text' in conversation_dict:
            return conversation_dict['text']
        elif 'message' in conversation_dict:
            return conversation_dict['message']
        elif 'episode_body' in conversation_dict:
            return conversation_dict['episode_body']
        else:
            # Fallback: use the whole dict as string representation
            return str(conversation_dict)
    
    def _filter_by_quality(self, items: List[Dict], category: str) -> List[Dict]:
        """Filter items by quality threshold and add relevance scores."""
        if not items:
            return []
        
        filtered_items = []
        for i, item in enumerate(items, 1):
            score = self._calculate_relevance_score(i, len(items))
            
            if score >= self.quality_threshold:
                item_with_score = item.copy()
                item_with_score['relevance_score'] = round(score, 3)
                item_with_score['rank'] = i
                item_with_score['memory_type'] = category
                
                # Extract and normalize the content based on category
                if category == 'fact':
                    item_with_score['fact'] = self._extract_fact_text(item)
                elif category == 'entity':
                    entity_info = self._extract_entity_info(item)
                    item_with_score['name'] = entity_info['name']
                    item_with_score['summary'] = entity_info['summary']
                elif category == 'conversation':
                    item_with_score['content'] = self._extract_conversation_content(item)
                
                filtered_items.append(item_with_score)
        
        return filtered_items
    
    async def get_similar_memories(self, 
                                 user_query: str, 
                                 max_results: Optional[int] = None,
                                 custom_threshold: Optional[float] = None) -> Dict[str, Any]:
        """
        Get similar memories for a user query.
        
        Args:
            user_query: The user's question or message
            max_results: Override default max results
            custom_threshold: Override default quality threshold
            
        Returns:
            Dictionary containing:
            - facts: List of relevant facts
            - entities: List of relevant entities (if enabled)
            - conversations: List of relevant conversations (if enabled)
            - total_memories: Total number of memories found
            - query: The original query
        """
        if not self.is_initialized:
            await self.initialize()
        
        max_results = max_results or self.default_max_results
        threshold = custom_threshold or self.quality_threshold
        
        try:
            # Get comprehensive search results
            results = await self.searcher.search_comprehensive(user_query, max_results)
            
            # Handle error case
            if 'error' in results:
                self.logger.error(f"Memory search error: {results['error']}")
                return {
                    'query': user_query,
                    'facts': [],
                    'entities': [],
                    'conversations': [],
                    'total_memories': 0,
                    'search_successful': False,
                    'error': results['error']
                }
            
            # Process and filter results
            memory_response = {
                'query': user_query,
                'facts': [],
                'entities': [],
                'conversations': [],
                'total_memories': 0,
                'search_successful': True
            }
            
            # Process facts (always included)
            if 'facts' in results and results['facts']:
                memory_response['facts'] = self._filter_by_quality(results['facts'], 'fact')
            
            # Process entities (if enabled)
            if self.include_entities and 'entities' in results and results['entities']:
                memory_response['entities'] = self._filter_by_quality(results['entities'], 'entity')
            
            # Process conversations (if enabled)
            if self.include_conversations and 'conversations' in results and results['conversations']:
                memory_response['conversations'] = self._filter_by_quality(results['conversations'], 'conversation')
            
            # Calculate total memories
            memory_response['total_memories'] = (
                len(memory_response['facts']) + 
                len(memory_response['entities']) + 
                len(memory_response['conversations'])
            )
            
            self.logger.debug(f"Found {memory_response['total_memories']} relevant memories for query: '{user_query}'")
            
            return memory_response
            
        except Exception as e:
            self.logger.error(f"Error searching memories for query '{user_query}': {e}")
            return {
                'query': user_query,
                'facts': [],
                'entities': [],
                'conversations': [],
                'total_memories': 0,
                'search_successful': False,
                'error': str(e)
            }
    
    async def get_facts_only(self, user_query: str, max_results: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        Get only relevant facts for a user query (simplified interface).
        
        Args:
            user_query: The user's question or message
            max_results: Override default max results
            
        Returns:
            List of relevant facts with scores
        """
        max_results = max_results or self.default_max_results
        
        if not self.is_initialized:
            await self.initialize()
        
        try:
            facts = await self.searcher.search_facts(user_query, max_results)
            return self._filter_by_quality(facts, 'fact')
        except Exception as e:
            self.logger.error(f"Error getting facts for query '{user_query}': {e}")
            return []
    
    def format_memories_for_agent(self, memories: Dict[str, Any], max_context_length: int = 2000) -> str:
        """
        Format memories into a text string suitable for passing to a chat agent.
        
        Args:
            memories: Result from get_similar_memories()
            max_context_length: Maximum length of formatted text
            
        Returns:
            Formatted string of relevant memories
        """
        if not memories.get('search_successful', False) or memories['total_memories'] == 0:
            return "No relevant memories found."
        
        context_parts = []
        
        # Add facts
        if memories['facts']:
            context_parts.append("RELEVANT FACTS:")
            for fact in memories['facts'][:5]:  # Limit to top 5 facts
                score = fact.get('relevance_score', 0)
                fact_text = fact.get('fact', '')
                context_parts.append(f"- [{score:.2f}] {fact_text}")
        
        # Add entities if available
        if memories['entities']:
            context_parts.append("\nRELEVANT ENTITIES:")
            for entity in memories['entities'][:3]:  # Limit to top 3 entities
                score = entity.get('relevance_score', 0)
                name = entity.get('name', '')
                summary = entity.get('summary', '')[:100] + "..." if len(entity.get('summary', '')) > 100 else entity.get('summary', '')
                context_parts.append(f"- [{score:.2f}] {name}: {summary}")
        
        # Add conversations if available
        if memories['conversations']:
            context_parts.append("\nRELEVANT PAST CONVERSATIONS:")
            for conv in memories['conversations'][:2]:  # Limit to top 2 conversations
                score = conv.get('relevance_score', 0)
                content = conv.get('content', '')[:150] + "..." if len(conv.get('content', '')) > 150 else conv.get('content', '')
                context_parts.append(f"- [{score:.2f}] {content}")
        
        formatted_context = "\n".join(context_parts)
        
        # Truncate if too long
        if len(formatted_context) > max_context_length:
            formatted_context = formatted_context[:max_context_length-3] + "..."
        
        return formatted_context


# Convenience function for quick integration
async def create_agent_memory_search(quality_threshold: float = 0.5) -> AgentMemorySearch:
    """
    Create and initialize an AgentMemorySearch instance.
    
    Args:
        quality_threshold: Minimum relevance score threshold
        
    Returns:
        Initialized AgentMemorySearch instance
    """
    memory_search = AgentMemorySearch(quality_threshold=quality_threshold)
    await memory_search.initialize()
    return memory_search


# Example usage
async def example_agent_integration():
    """Example of how to integrate with a chat agent."""
    
    # Initialize once at startup
    memory_search = await create_agent_memory_search(quality_threshold=0.6)
    
    # Simulate user queries
    user_queries = [
        "What are the master's responsibilities?",
        "Tell me about Rohan",
        "How do I replace the GP-150 battery?"
    ]
    
    for query in user_queries:
        print(f"\n🔍 User Query: '{query}'")
        print("-" * 50)
        
        # Get memories
        memories = await memory_search.get_similar_memories(query)
        
        # Format for agent
        context = memory_search.format_memories_for_agent(memories)
        print("Context for Agent:")
        print(context)
        print()


if __name__ == "__main__":
    asyncio.run(example_agent_integration()) 