#!/usr/bin/env python3
"""
Memory Search Module - Graphiti Implementation

This implements memory search using Graphiti core for graph-based memory storage and retrieval.
"""

import os
import json
import asyncio
from datetime import datetime, timezone
from typing import List, Dict, Any, Optional

from graphiti_core.graphiti import Graphiti
from graphiti_core.search.search_config_recipes import COMBINED_HYBRID_SEARCH_RRF
from graphiti_core.llm_client import LLMConfig, OpenAIClient


class MemorySearcher:
    """A class for searching memory using Graphiti client."""
    
    def __init__(self):
        self.graphiti_client: Optional[Graphiti] = None
        self.neo4j_uri = os.getenv("NEO4J_URI", "neo4j://127.0.0.1:7687")
        self.neo4j_user = os.getenv("NEO4J_USER", "neo4j")
        self.neo4j_password = os.getenv("NEO4J_PASSWORD")
        self.openai_api_key = os.getenv("OPENAI_API_KEY")
        self.model_name = os.getenv("MODEL_NAME", "gpt-4o")
        
        if not self.neo4j_password:
            raise ValueError("NEO4J_PASSWORD environment variable is required")
        if not self.openai_api_key:
            raise ValueError("OPENAI_API_KEY environment variable is required")
    
    async def _ensure_client_initialized(self):
        """Ensure the Graphiti client is initialized."""
        if self.graphiti_client is None:
            # Configure LLM client
            llm_config = LLMConfig(
                model=self.model_name,
                temperature=0.0,
                max_tokens=4000
            )
            llm_client = OpenAIClient(llm_config)
            
            # Initialize Graphiti client
            self.graphiti_client = Graphiti(
                uri=self.neo4j_uri,
                user=self.neo4j_user,
                password=self.neo4j_password,
                llm_client=llm_client
            )
            
            # Ensure indices and constraints are built
            await self.graphiti_client.build_indices_and_constraints()
    
    async def search_facts(
        self, 
        query: str, 
        max_results: int = 10,
        group_ids: Optional[List[str]] = None,
        center_around: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Search for facts related to a given query.
        
        Args:
            query: The search query string
            max_results: Maximum number of results to return (default: 10)
            group_ids: Optional list of group IDs to filter by
            center_around: Optional UUID to center the search around
            
        Returns:
            List of facts/edges related to the query
        """
        try:
            # Ensure client is initialized
            await self._ensure_client_initialized()
            assert self.graphiti_client is not None, "Graphiti client should be initialized"
            
            # Configure search
            config = COMBINED_HYBRID_SEARCH_RRF.model_copy(deep=True)
            config.limit = max_results
            
            # Perform search
            results = await self.graphiti_client.search_(
                query=query,
                config=config,
                group_ids=group_ids or [],
                center_node_uuid=center_around
            )
            
            # Extract and format facts (edges)
            facts = []
            if hasattr(results, 'edges') and results.edges:
                for edge in results.edges:
                    fact_dict = edge.model_dump(mode='json', exclude={'fact_embedding'})
                    # Remove embedding from attributes if present
                    if 'attributes' in fact_dict and isinstance(fact_dict['attributes'], dict):
                        fact_dict['attributes'].pop('fact_embedding', None)
                    facts.append(fact_dict)
            
            return facts
            
        except Exception as e:
            print(f"Error searching facts: {e}")
            return []
    
    async def search_entities(self, query: str, max_results: int = 10) -> List[Dict[str, Any]]:
        """
        Search for entities only.
        
        Args:
            query: The search query
            max_results: Maximum number of results to return
            
        Returns:
            List of entity dictionaries
        """
        try:
            # Get comprehensive results and extract entities
            comprehensive_results = await self.search_comprehensive(query, max_results)
            return comprehensive_results.get('entities', [])
        except Exception as e:
            print(f"Error searching entities: {e}")
            return []
    
    async def search_conversations(self, query: str, max_results: int = 10) -> List[Dict[str, Any]]:
        """
        Search for conversations only.
        
        Args:
            query: The search query
            max_results: Maximum number of results to return
            
        Returns:
            List of conversation dictionaries
        """
        try:
            # Get comprehensive results and extract conversations
            comprehensive_results = await self.search_comprehensive(query, max_results)
            return comprehensive_results.get('conversations', [])
        except Exception as e:
            print(f"Error searching conversations: {e}")
            return []
    
    async def search_comprehensive(
        self, 
        query: str, 
        max_results: int = 10,
        group_ids: Optional[List[str]] = None,
        center_around: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Comprehensive search that returns facts, entities, and conversations.
        
        Args:
            query: The search query string
            max_results: Maximum number of results to return (default: 10)
            group_ids: Optional list of group IDs to filter by
            center_around: Optional UUID to center the search around
            
        Returns:
            Dictionary containing facts, entities, and conversations
        """
        try:
            # Ensure client is initialized
            await self._ensure_client_initialized()
            assert self.graphiti_client is not None, "Graphiti client should be initialized"
            
            # Configure search
            config = COMBINED_HYBRID_SEARCH_RRF.model_copy(deep=True)
            config.limit = max_results
            
            # Perform search
            results = await self.graphiti_client.search_(
                query=query,
                config=config,
                group_ids=group_ids or [],
                center_node_uuid=center_around
            )
            
            output = {}
            
            # Process facts (edges)
            if hasattr(results, 'edges') and results.edges:
                output['facts'] = []
                for edge in results.edges:
                    edge_dict = edge.model_dump(mode='json', exclude={'fact_embedding'})
                    if 'attributes' in edge_dict and isinstance(edge_dict['attributes'], dict):
                        edge_dict['attributes'].pop('fact_embedding', None)
                    output['facts'].append(edge_dict)
            
            # Process entities (nodes)
            if hasattr(results, 'nodes') and results.nodes:
                output['entities'] = []
                for node in results.nodes:
                    node_dict = node.model_dump(mode='json', exclude={'node_embedding', 'name_embedding'})
                    if 'attributes' in node_dict and isinstance(node_dict['attributes'], dict):
                        node_dict['attributes'].pop('node_embedding', None)
                        node_dict['attributes'].pop('name_embedding', None)
                    output['entities'].append(node_dict)
            
            # Process conversations (episodes)
            if hasattr(results, 'episodes') and results.episodes:
                output['conversations'] = []
                for episode in results.episodes:
                    episode_dict = episode.model_dump(mode='json', exclude={'episode_embedding'})
                    output['conversations'].append(episode_dict)
            
            return output
            
        except Exception as e:
            print(f"Error in comprehensive search: {e}")
            return {"error": str(e)}


# Convenience functions for easy usage
async def search_memory_facts(query: str, max_results: int = 10) -> List[Dict[str, Any]]:
    """
    Simple function to search for facts related to a query.
    
    Args:
        query: The search query string
        max_results: Maximum number of results to return (default: 10)
        
    Returns:
        List of facts related to the query
    """
    searcher = MemorySearcher()
    return await searcher.search_facts(query, max_results)


async def search_memory_comprehensive(query: str, max_results: int = 10) -> Dict[str, Any]:
    """
    Simple function to perform comprehensive memory search.
    
    Args:
        query: The search query string
        max_results: Maximum number of results to return (default: 10)
        
    Returns:
        Dictionary containing facts, entities, and conversations
    """
    searcher = MemorySearcher()
    return await searcher.search_comprehensive(query, max_results)


# Example usage
if __name__ == "__main__":
    async def main():
        # Example 1: Search for facts only
        print("Searching for facts about 'user authentication'...")
        facts = await search_memory_facts("user authentication", max_results=5)
        print(f"Found {len(facts)} facts:")
        print(json.dumps(facts, indent=2))
        
        print("\n" + "="*50 + "\n")
        
        # Example 2: Comprehensive search
        print("Comprehensive search for 'project status'...")
        results = await search_memory_comprehensive("project status", max_results=5)
        print(json.dumps(results, indent=2))
    
    # Run the example
    asyncio.run(main()) 