"""Memory management system using HawkinsDB."""
from typing import Dict, Any, List, Optional
from hawkinsdb import HawkinsDB
import logging
from datetime import datetime

logger = logging.getLogger(__name__)

class MemoryManager:
    """Manages agent memory using HawkinsDB for persistent storage"""

    def __init__(self, db: HawkinsDB):
        self.db = db

    def store_interaction(self, task: str, result: Dict[str, Any]):
        """Store an agent interaction in episodic memory"""
        try:
            memory_data = {
                "name": f"task_{datetime.now().isoformat()}",
                "column": "Episodic",
                "properties": {
                    "task": task,
                    "timestamp": datetime.now().isoformat(),
                    "thoughts": result.get('thoughts', []),
                    "actions": result.get('actions', []),
                    "observations": result.get('observations', []),
                    "final_answer": result.get('final_answer'),
                    "context": result.get('context', ''),
                    "sources": result.get('sources', []),
                    "type": "interaction"
                }
            }

            self.db.add_entity(memory_data)
            logger.info(f"Stored interaction memory for task: {task}")
        except Exception as e:
            logger.error(f"Error storing interaction memory: {str(e)}")

    def store_knowledge(self, topic: str, content: Dict[str, Any]):
        """Store semantic knowledge"""
        try:
            knowledge_data = {
                "name": f"knowledge_{topic}",
                "column": "Semantic",
                "properties": {
                    "topic": topic,
                    "content": content,
                    "timestamp": datetime.now().isoformat(),
                    "type": "knowledge",
                    "sources": content.get('sources', []),
                    "context": content.get('context', '')
                }
            }

            self.db.add_entity(knowledge_data)
            logger.info(f"Stored semantic knowledge for topic: {topic}")
        except Exception as e:
            logger.error(f"Error storing semantic knowledge: {str(e)}")

    def store_skill(self, skill_name: str, steps: List[str]):
        """Store a procedural memory (skill)"""
        try:
            skill_data = {
                "name": skill_name,
                "column": "Procedural",
                "properties": {
                    "steps": steps,
                    "timestamp": datetime.now().isoformat(),
                    "type": "skill"
                }
            }
            self.db.add_entity(skill_data)
            logger.info(f"Stored procedural memory for skill: {skill_name}")
        except Exception as e:
            logger.error(f"Error storing procedural memory: {str(e)}")

    def get_relevant_context(self, query: str) -> Dict[str, Any]:
        """Get relevant context from memory"""
        try:
            # Query all memory types with semantic search
            memories = list(self.db.query_frames(query))

            # Filter and organize memories by type
            episodic = []
            semantic = []

            # Process and organize memories
            for memory in memories:
                memory_type = memory.get("properties", {}).get("type")
                if memory_type == "interaction":
                    if len(episodic) < 5:  # Limit to most relevant 5
                        episodic.append(memory)
                elif memory_type == "knowledge":
                    if len(semantic) < 5:  # Limit to most relevant 5
                        semantic.append(memory)

            return {
                "episodic_memories": episodic,
                "semantic_memories": semantic,
                "all_memories": memories[:10]  # Limit total memories
            }
        except Exception as e:
            logger.error(f"Error retrieving context: {str(e)}")
            return {
                "episodic_memories": [],
                "semantic_memories": [],
                "all_memories": []
            }

    def get_recent_interactions(self, limit: int = 5) -> List[Dict[str, Any]]:
        """Get recent agent interactions"""
        try:
            # Query and limit the results
            interactions = list(self.db.query_frames("type:interaction"))
            return interactions[:min(limit, len(interactions))]
        except Exception as e:
            logger.error(f"Error retrieving recent interactions: {str(e)}")
            return []

    def get_skill(self, skill_name: str) -> Optional[Dict[str, Any]]:
        """Retrieve a procedural memory (skill)"""
        try:
            skills = list(self.db.query_frames(f"name:{skill_name} AND column:Procedural"))
            return skills[0] if skills else None
        except Exception as e:
            logger.error(f"Error retrieving skill: {str(e)}")
            return None