"""
Memory Manager Implementation

This module handles both short-term conversational memory and long-term
knowledge storage using vector embeddings for retrieval-augmented generation.
"""

import os
import json
import uuid
import time
from typing import List, Dict, Any, Optional
from dataclasses import dataclass, asdict
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer


@dataclass
class MemoryItem:
    """Represents a single memory item"""
    id: str
    content: str
    memory_type: str  # 'interaction', 'knowledge', 'observation'
    timestamp: float
    metadata: Dict[str, Any]
    embedding: Optional[List[float]] = None


class MemoryManager:
    """
    Manages agent memory using vector embeddings for semantic retrieval
    """
    
    def __init__(self, persist_path: str = "./agent_memory"):
        """
        Initialize the memory manager
        
        Args:
            persist_path: Directory to persist memory data
        """
        self.persist_path = persist_path
        os.makedirs(persist_path, exist_ok=True)
        
        # Initialize embedding model
        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        
        # Initialize ChromaDB for vector storage
        self.chroma_client = chromadb.PersistentClient(
            path=os.path.join(persist_path, "chroma_db"),
            settings=Settings(anonymized_telemetry=False)
        )
        
        # Create collections for different types of memory
        self.interactions_collection = self._get_or_create_collection("interactions")
        self.knowledge_collection = self._get_or_create_collection("knowledge")
        self.observations_collection = self._get_or_create_collection("observations")
        
        # In-memory cache for recent interactions
        self.recent_interactions = []
        self.max_recent_interactions = 10

    def _get_or_create_collection(self, name: str):
        """Get or create a ChromaDB collection"""
        try:
            return self.chroma_client.get_collection(name)
        except:
            return self.chroma_client.create_collection(name)

    def _generate_embedding(self, text: str) -> List[float]:
        """Generate embedding for text"""
        return self.embedding_model.encode(text).tolist()

    def store_interaction(self, interaction_data: Dict[str, Any]) -> str:
        """
        Store a user-agent interaction
        
        Args:
            interaction_data: Dictionary containing query, response, tool_outputs, etc.
            
        Returns:
            Unique ID of the stored interaction
        """
        interaction_id = str(uuid.uuid4())
        
        # Create searchable content from interaction
        searchable_content = f"""
        Query: {interaction_data.get('query', '')}
        Response: {interaction_data.get('response', '')}
        Tools Used: {', '.join(interaction_data.get('tool_outputs', {}).keys())}
        """
        
        # Generate embedding
        embedding = self._generate_embedding(searchable_content)
        
        # Store in ChromaDB
        self.interactions_collection.add(
            ids=[interaction_id],
            documents=[searchable_content],
            embeddings=[embedding],
            metadatas=[{
                "timestamp": interaction_data.get('timestamp', time.time()),
                "query": interaction_data.get('query', ''),
                "response_length": len(interaction_data.get('response', '')),
                "tools_used": ", ".join(interaction_data.get('tool_outputs', {}).keys())
            }]
        )
        
        # Add to recent interactions cache
        memory_item = MemoryItem(
            id=interaction_id,
            content=searchable_content,
            memory_type="interaction",
            timestamp=interaction_data.get('timestamp', time.time()),
            metadata=interaction_data,
            embedding=embedding
        )
        
        self.recent_interactions.append(memory_item)
        if len(self.recent_interactions) > self.max_recent_interactions:
            self.recent_interactions.pop(0)
        
        return interaction_id

    def store_knowledge(self, content: str, source: str, metadata: Optional[Dict] = None) -> str:
        """
        Store knowledge/information for future retrieval
        
        Args:
            content: The knowledge content to store
            source: Source of the knowledge (e.g., 'web_scraping', 'user_input')
            metadata: Additional metadata
            
        Returns:
            Unique ID of the stored knowledge
        """
        knowledge_id = str(uuid.uuid4())
        
        # Generate embedding
        embedding = self._generate_embedding(content)
        
        # Prepare metadata
        full_metadata = {
            "source": source,
            "timestamp": time.time(),
            "content_length": len(content)
        }
        if metadata:
            full_metadata.update(metadata)
        
        # Store in ChromaDB
        self.knowledge_collection.add(
            ids=[knowledge_id],
            documents=[content],
            embeddings=[embedding],
            metadatas=[full_metadata]
        )
        
        return knowledge_id

    def store_observation(self, observation: str, context: Dict[str, Any]) -> str:
        """
        Store an observation or insight
        
        Args:
            observation: The observation text
            context: Context in which the observation was made
            
        Returns:
            Unique ID of the stored observation
        """
        observation_id = str(uuid.uuid4())
        
        # Generate embedding
        embedding = self._generate_embedding(observation)
        
        # Store in ChromaDB
        self.observations_collection.add(
            ids=[observation_id],
            documents=[observation],
            embeddings=[embedding],
            metadatas=[{
                "timestamp": time.time(),
                "context": json.dumps(context),
                "observation_length": len(observation)
            }]
        )
        
        return observation_id

    def retrieve_relevant_context(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
        """
        Retrieve relevant context for a query using semantic search
        
        Args:
            query: The query to find relevant context for
            max_results: Maximum number of results to return
            
        Returns:
            List of relevant memory items
        """
        query_embedding = self._generate_embedding(query)
        relevant_items = []
        
        # Search interactions
        try:
            interaction_results = self.interactions_collection.query(
                query_embeddings=[query_embedding],
                n_results=min(max_results, 3)
            )
            
            for i, doc in enumerate(interaction_results['documents'][0]):
                relevant_items.append({
                    "type": "interaction",
                    "content": doc,
                    "metadata": interaction_results['metadatas'][0][i],
                    "distance": interaction_results['distances'][0][i] if 'distances' in interaction_results else 0
                })
        except Exception as e:
            print(f"Error searching interactions: {e}")
        
        # Search knowledge
        try:
            knowledge_results = self.knowledge_collection.query(
                query_embeddings=[query_embedding],
                n_results=min(max_results, 3)
            )
            
            for i, doc in enumerate(knowledge_results['documents'][0]):
                relevant_items.append({
                    "type": "knowledge",
                    "content": doc,
                    "metadata": knowledge_results['metadatas'][0][i],
                    "distance": knowledge_results['distances'][0][i] if 'distances' in knowledge_results else 0
                })
        except Exception as e:
            print(f"Error searching knowledge: {e}")
        
        # Search observations
        try:
            observation_results = self.observations_collection.query(
                query_embeddings=[query_embedding],
                n_results=min(max_results, 2)
            )
            
            for i, doc in enumerate(observation_results['documents'][0]):
                relevant_items.append({
                    "type": "observation",
                    "content": doc,
                    "metadata": observation_results['metadatas'][0][i],
                    "distance": observation_results['distances'][0][i] if 'distances' in observation_results else 0
                })
        except Exception as e:
            print(f"Error searching observations: {e}")
        
        # Sort by relevance (lower distance = more relevant)
        relevant_items.sort(key=lambda x: x.get('distance', 1.0))
        
        return relevant_items[:max_results]

    def get_recent_interactions(self, count: int = 5) -> List[MemoryItem]:
        """Get the most recent interactions"""
        return self.recent_interactions[-count:]

    def search_by_keyword(self, keyword: str, memory_type: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        Search memory by keyword
        
        Args:
            keyword: Keyword to search for
            memory_type: Type of memory to search ('interaction', 'knowledge', 'observation')
            
        Returns:
            List of matching memory items
        """
        results = []
        collections_to_search = []
        
        if memory_type is None or memory_type == "interaction":
            collections_to_search.append(("interaction", self.interactions_collection))
        if memory_type is None or memory_type == "knowledge":
            collections_to_search.append(("knowledge", self.knowledge_collection))
        if memory_type is None or memory_type == "observation":
            collections_to_search.append(("observation", self.observations_collection))
        
        for mem_type, collection in collections_to_search:
            try:
                # ChromaDB doesn't have direct keyword search, so we use semantic search
                keyword_embedding = self._generate_embedding(keyword)
                search_results = collection.query(
                    query_embeddings=[keyword_embedding],
                    n_results=10
                )
                
                for i, doc in enumerate(search_results['documents'][0]):
                    if keyword.lower() in doc.lower():
                        results.append({
                            "type": mem_type,
                            "content": doc,
                            "metadata": search_results['metadatas'][0][i]
                        })
            except Exception as e:
                print(f"Error searching {mem_type}: {e}")
        
        return results

    def get_stats(self) -> Dict[str, Any]:
        """Get memory statistics"""
        stats = {
            "recent_interactions": len(self.recent_interactions),
            "total_interactions": 0,
            "total_knowledge": 0,
            "total_observations": 0
        }
        
        try:
            stats["total_interactions"] = self.interactions_collection.count()
        except:
            pass
        
        try:
            stats["total_knowledge"] = self.knowledge_collection.count()
        except:
            pass
        
        try:
            stats["total_observations"] = self.observations_collection.count()
        except:
            pass
        
        return stats

    def clear_all_memory(self):
        """Clear all stored memory (use with caution)"""
        try:
            self.chroma_client.delete_collection("interactions")
            self.chroma_client.delete_collection("knowledge")
            self.chroma_client.delete_collection("observations")
            
            # Recreate collections
            self.interactions_collection = self._get_or_create_collection("interactions")
            self.knowledge_collection = self._get_or_create_collection("knowledge")
            self.observations_collection = self._get_or_create_collection("observations")
            
            # Clear recent interactions
            self.recent_interactions = []
            
        except Exception as e:
            print(f"Error clearing memory: {e}")

    def export_memory(self, output_path: str):
        """Export all memory to a JSON file"""
        export_data = {
            "interactions": [],
            "knowledge": [],
            "observations": [],
            "export_timestamp": time.time()
        }
        
        # Export interactions
        try:
            all_interactions = self.interactions_collection.get()
            for i, doc in enumerate(all_interactions['documents']):
                export_data["interactions"].append({
                    "id": all_interactions['ids'][i],
                    "content": doc,
                    "metadata": all_interactions['metadatas'][i]
                })
        except Exception as e:
            print(f"Error exporting interactions: {e}")
        
        # Export knowledge
        try:
            all_knowledge = self.knowledge_collection.get()
            for i, doc in enumerate(all_knowledge['documents']):
                export_data["knowledge"].append({
                    "id": all_knowledge['ids'][i],
                    "content": doc,
                    "metadata": all_knowledge['metadatas'][i]
                })
        except Exception as e:
            print(f"Error exporting knowledge: {e}")
        
        # Export observations
        try:
            all_observations = self.observations_collection.get()
            for i, doc in enumerate(all_observations['documents']):
                export_data["observations"].append({
                    "id": all_observations['ids'][i],
                    "content": doc,
                    "metadata": all_observations['metadatas'][i]
                })
        except Exception as e:
            print(f"Error exporting observations: {e}")
        
        # Write to file
        with open(output_path, 'w') as f:
            json.dump(export_data, f, indent=2)
        
        return export_data

