"""Core Agent implementation with memory and tool management."""
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from .tools.base import Tool
from .memory.manager import MemoryManager
from hawkinsdb import HawkinsDB
from hawkins_rag import HawkinsRAG
import logging
import os
import tempfile
from datetime import datetime

logger = logging.getLogger(__name__)

@dataclass
class AgentConfig:
    """Configuration for an AI agent"""
    name: str
    role: str
    goal: str
    tools: Optional[List[Tool]] = None
    model: str = "gpt-4"
    temperature: float = 0.7
    db_path: str = "agent_memory.db"
    knowledge_base_path: str = "knowledge.db"

class Agent:
    def __init__(self, config: AgentConfig):
        """Initialize an AI agent with the given configuration"""
        self.config = config
        self.db = HawkinsDB(storage_type='sqlite', db_path=config.db_path)
        self.memory = MemoryManager(self.db)
        # Initialize RAG with proper configuration
        self.rag = HawkinsRAG(config={
            "storage_type": "sqlite",
            "db_path": config.knowledge_base_path,
            "chunk_size": 500  # Smaller chunks for better retrieval
        })
        self.tools = {tool.name: tool for tool in (config.tools or [])}
        self.history = []

    async def load_knowledge(self, content: str, source_type: str = "text", metadata: Dict[str, Any] = None) -> Dict[str, Any]:
        """Load content into the knowledge base"""
        try:
            logger.info(f"Loading {source_type} content into knowledge base")
            metadata = metadata or {}

            # Write content to temporary file for text documents
            if source_type == "text":
                with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as temp_file:
                    temp_file.write(content.strip())
                    temp_file.flush()

                    # Load document into RAG system
                    result = self.rag.load_document(
                        temp_file.name,
                        source_type=source_type
                    )

                    # Clean up temp file
                    os.unlink(temp_file.name)
            else:
                # Handle other document types directly
                result = self.rag.load_document(content, source_type=source_type)

            # Store metadata and document info in memory
            if result:
                topic = metadata.get("topic", "document")
                logger.info(f"Storing metadata for topic: {topic}")
                self.memory.store_knowledge(
                    topic=topic,
                    content={
                        "type": source_type,
                        "metadata": metadata,
                        "summary": "Document loaded into knowledge base",
                        "status": "success"
                    }
                )

                return {
                    "success": True,
                    "message": "Content loaded successfully",
                    "result": result
                }

            return {
                "success": False,
                "error": "Failed to load document into knowledge base"
            }

        except Exception as e:
            logger.error(f"Error loading knowledge: {str(e)}")
            return {
                "success": False,
                "error": str(e)
            }

    async def query_knowledge(self, query: str) -> Dict[str, Any]:
        """Query the knowledge base with enhanced context handling"""
        try:
            logger.info(f"Querying knowledge base: {query}")

            # Get relevant memories before querying
            memory_context = self.memory.get_relevant_context(query)
            related_answers = [
                mem['properties'].get('final_answer')
                for mem in memory_context['episodic_memories']
                if mem['properties'].get('final_answer')
            ]

            # Enhance query with related context if available
            enhanced_query = query
            if related_answers:
                context_str = "\n".join(f"Previous relevant answer: {answer}" for answer in related_answers[:2])
                enhanced_query = f"{query}\nContext:\n{context_str}"

            # Query the RAG system with enhanced query
            response = self.rag.query(enhanced_query)

            # Process and normalize the response
            if isinstance(response, dict):
                answer = response.get('response', response)
                sources = response.get('sources', [])
                context = response.get('context', '')

                if isinstance(answer, dict):
                    answer = answer.get('response', "No specific answer found.")
            else:
                answer = str(response)
                sources = []
                context = ''

            # Store interaction with enhanced context
            self.memory.store_interaction(
                task=query,
                result={
                    "thoughts": ["Used knowledge base for query"],
                    "query": query,
                    "final_answer": answer,
                    "sources": sources,
                    "context": context,
                    "related_memories": related_answers,
                    "timestamp": datetime.now().isoformat()
                }
            )

            return {
                "success": True,
                "answer": answer,
                "sources": sources,
                "context": context,
                "related_memories": related_answers,
                "message": "Query processed successfully"
            }

        except Exception as e:
            logger.error(f"Error querying knowledge base: {str(e)}")
            return {
                "success": False,
                "error": str(e)
            }

    async def execute(self, task: str, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Execute a task using available tools and memory"""
        context = context or {}
        logger.info(f"Executing task: {task}")

        try:
            # Add memory context
            memory_context = self.memory.get_relevant_context(task)
            context.update(memory_context)

            # Check if this is a knowledge base operation
            if context.get('knowledge_action'):
                action = context['knowledge_action']
                if action['operation'] == 'load':
                    return await self.load_knowledge(
                        content=action['content'],
                        source_type=action.get('source_type', 'text'),
                        metadata=action.get('metadata', {})
                    )
                elif action['operation'] == 'query':
                    return await self.query_knowledge(action.get('query', task))

            # Enhanced query routing logic
            # Check if the task looks like a knowledge query
            is_query = any(keyword in task.lower() for keyword in [
                'what', 'how', 'why', 'when', 'who', 'where', 
                'explain', 'describe', 'tell me about', 'find',
                'search', 'look up'
            ])

            if is_query:
                # Route to knowledge base query
                return await self.query_knowledge(task)

            # Handle tool-specific tasks
            tool_result = await self._execute_tool(task, context)
            if tool_result:
                return tool_result

            # Store the interaction
            self.memory.store_interaction(
                task=task,
                result={
                    'thoughts': ['Task completed without specific tool or knowledge query'],
                    'final_answer': 'Task executed successfully',
                    'context': context,
                    'timestamp': datetime.now().isoformat()
                }
            )

            return {
                'success': True,
                'message': 'Task completed successfully',
                'context': context
            }

        except Exception as e:
            logger.error(f"Error executing task: {str(e)}")
            return {
                'success': False,
                'error': f"Error executing task: {str(e)}"
            }

    async def _execute_tool(self, task: str, context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Execute a specific tool based on the task"""
        try:
            # Check for calculator tasks
            if "calculate" in task.lower():
                tool = self.tools.get('calculator')
                if tool:
                    result = await tool.execute(task)
                    if result['success']:
                        self.memory.store_interaction(
                            task=task,
                            result={
                                'thoughts': ['Used calculator tool'],
                                'final_answer': f"The result is {result['result']}"
                            }
                        )
                        return {
                            'success': True,
                            'answer': f"The result is {result['result']}",
                            'result': result
                        }

            return None

        except Exception as e:
            logger.error(f"Error executing tool: {str(e)}")
            return {
                'success': False,
                'error': f"Error executing tool: {str(e)}"
            }

    def add_tool(self, tool: Tool):
        """Add a new tool to the agent's capabilities"""
        self.tools[tool.name] = tool

    def remove_tool(self, tool_name: str):
        """Remove a tool from the agent's capabilities"""
        if tool_name in self.tools:
            del self.tools[tool_name]