"""
SingleAgent class for Metis Agent.

This module provides the main agent class that orchestrates all components.
"""
import os
import time
from typing import Dict, List, Any, Optional, Union
import json

from ..core.intent_router import IntentRouter
from ..core.planner import Planner
from ..core.task_manager import TaskManager
from ..core.scheduler import Scheduler
from ..core.llm_interface import get_llm
from ..tools.registry import initialize_tools
from ..memory.sqlite_store import SQLiteMemory
from ..memory.titans.titans_adapter import TitansMemoryAdapter


class SingleAgent:
    """
    Main agent class that orchestrates all components.
    """
    
    def __init__(
        self,
        use_titans_memory: bool = False,
        tools: Optional[List[Any]] = None,
        llm_provider: Optional[str] = None,
        llm_model: Optional[str] = None,
        memory_path: Optional[str] = None,
        task_file: Optional[str] = None
    ):
        """
        Initialize the agent.
        
        Args:
            use_titans_memory: Whether to use Titans memory
            tools: List of tool instances to use (if None, uses all available tools)
            llm_provider: LLM provider to use (if None, auto-detects based on available API keys)
            llm_model: LLM model to use
            memory_path: Path to memory database
            task_file: Path to task file
        """
        # Set up LLM
        from ..core.llm_interface import configure_llm
        if llm_provider:
            # User specified a provider
            if llm_model:
                configure_llm(llm_provider, llm_model)
            else:
                configure_llm(llm_provider)
        # If no provider specified, get_llm() will auto-detect based on available API keys
            
        self.llm = get_llm()
        
        # Set up components
        self.intent_router = IntentRouter()
        
        # Set up task management
        if task_file is None:
            self.task_file = os.path.join(os.getcwd(), "tasks.md")
        else:
            self.task_file = task_file
            
        self.planner = Planner(self.task_file)
        self.task_manager = TaskManager(self.task_file)
        self.scheduler = Scheduler()
        
        # Set up memory
        if memory_path is None:
            memory_dir = os.path.join(os.getcwd(), "memory")
            os.makedirs(memory_dir, exist_ok=True)
            self.memory_path = os.path.join(memory_dir, "memory.db")
        else:
            self.memory_path = memory_path
            
        self.memory = SQLiteMemory(self.memory_path)
        
        # Set up Titans memory if enabled
        self.titans_memory_enabled = use_titans_memory
        self.titans_adapter = None
        
        if self.titans_memory_enabled:
            try:
                self.titans_adapter = TitansMemoryAdapter(self)
                print("+ Titans memory initialized")
            except Exception as e:
                print(f"- Error initializing Titans memory: {e}")
                self.titans_memory_enabled = False
        
        # Set up tools
        if tools is None:
            self.tools = initialize_tools()
        else:
            self.tools = {tool.__class__.__name__: tool for tool in tools}
            
        # Set up agent identity
        self.agent_creation_date = time.strftime("%Y-%m-%d")
        self.agent_name = "Metis Agent"
        self.version = "0.1.3"
        
        print(f"+ {self.agent_name} v{self.version} initialized")
        
    def process_query(
        self,
        query: str,
        session_id: Optional[str] = None,
        tool_name: Optional[str] = None
    ) -> Union[str, Dict[str, Any]]:
        """
        Process a user query.
        
        Args:
            query: User query
            session_id: Session identifier
            tool_name: Name of tool to use (if None, selects automatically)
            
        Returns:
            Response as string or dictionary
        """
        if not query.strip():
            return "Query cannot be empty."
            
        # Generate session ID if not provided
        if session_id is None:
            session_id = f"session_{int(time.time())}"
            
        # Store query in memory
        self.memory.store_input(session_id, query)
        
        # Enhance query with Titans memory if enabled
        enhanced_context = ""
        if self.titans_memory_enabled and self.titans_adapter:
            try:
                enhancement = self.titans_adapter.enhance_query_processing(query, session_id)
                enhanced_context = enhancement.get("enhanced_context", "")
                if enhanced_context:
                    print(f"+ Enhanced query with {len(enhancement.get('relevant_memories', []))} memories")
            except Exception as e:
                print(f"- Error enhancing query with Titans memory: {e}")
        
        # Classify intent
        intent = self.intent_router.classify(query)
        print(f"+ Query classified as: {intent}")
        
        if intent == "question":
            # Handle as a question
            return self._handle_question(query, enhanced_context, session_id)
        else:
            # Handle as a task
            return self._handle_task(query, enhanced_context, session_id, tool_name)
            
    def _handle_question(
        self,
        query: str,
        enhanced_context: str,
        session_id: str
    ) -> str:
        """
        Handle a question by sending it directly to the LLM.
        
        Args:
            query: User query
            enhanced_context: Enhanced context from Titans memory
            session_id: Session identifier
            
        Returns:
            Answer to the question
        """
        # Get conversation context
        conversation_context = self.memory.get_context(session_id)
        
        # Create prompt with context
        prompt = [
            {"role": "system", "content": f"""You are {self.agent_name}, an advanced AI assistant.
Answer the user's question accurately, concisely, and helpfully.
If you don't know the answer, say so rather than making something up.

{enhanced_context}

Conversation history:
{conversation_context}"""},
            {"role": "user", "content": query}
        ]
        
        # Get response from LLM
        response = self.llm.chat(prompt)
        
        # Store response in memory
        self.memory.store_output(session_id, response)
        
        # Store in Titans memory if enabled
        if self.titans_memory_enabled and self.titans_adapter:
            try:
                self.titans_adapter.store_response(query, response, session_id)
            except Exception as e:
                print(f"⚠️ Error storing response in Titans memory: {e}")
        
        return response
        
    def _handle_task(
        self,
        query: str,
        enhanced_context: str,
        session_id: str,
        tool_name: Optional[str] = None
    ) -> Union[str, Dict[str, Any]]:
        """
        Handle a task by planning, scheduling, and executing it.
        
        Args:
            query: User query
            enhanced_context: Enhanced context from Titans memory
            session_id: Session identifier
            tool_name: Name of tool to use (if None, selects automatically)
            
        Returns:
            Task result
        """
        # Start new task session
        self.task_manager.start_new_task_within_session(query)
        
        # Create plan
        tasks = self.planner.create_plan(query)
        print(f"+ Created plan with {len(tasks)} tasks")
        
        # Add tasks to task manager
        self.task_manager.add_tasks(tasks, original_query=query)
        
        # Prioritize tasks
        prioritized_tasks = self.scheduler.prioritize_tasks(tasks)
        
        # Execute tasks
        results = []
        
        for task in prioritized_tasks:
            print(f"+ Executing task: {task}")
            
            # Select tool for task
            selected_tool = None
            
            if tool_name:
                # Use specified tool
                selected_tool = self.tools.get(tool_name)
                if not selected_tool:
                    print(f"- Tool '{tool_name}' not found, selecting automatically")
            
            if not selected_tool:
                # Select tool automatically
                for tool in self.tools.values():
                    if tool.can_handle(task):
                        selected_tool = tool
                        break
            
            if selected_tool:
                print(f"+ Using tool: {selected_tool.__class__.__name__}")
                
                try:
                    # Execute task with tool
                    task_result = selected_tool.execute(task)
                    results.append({
                        "task": task,
                        "tool": selected_tool.__class__.__name__,
                        "result": task_result
                    })
                    
                    # Mark task as complete
                    self.task_manager.mark_complete(task)
                    
                except Exception as e:
                    print(f"- Error executing task: {e}")
                    results.append({
                        "task": task,
                        "tool": selected_tool.__class__.__name__ if selected_tool else "None",
                        "error": str(e)
                    })
            else:
                print(f"- No suitable tool found for task: {task}")
                
                # Use LLM as fallback
                print("+ Using LLM as fallback")
                
                # Get conversation context
                conversation_context = self.memory.get_context(session_id)
                
                # Create prompt with context
                prompt = [
                    {"role": "system", "content": f"""You are {self.agent_name}, an advanced AI assistant.
Complete the following task to the best of your ability.

{enhanced_context}

Conversation history:
{conversation_context}"""},
                    {"role": "user", "content": task}
                ]
                
                # Get response from LLM
                task_result = self.llm.chat(prompt)
                
                results.append({
                    "task": task,
                    "tool": "LLM Fallback",
                    "result": task_result
                })
                
                # Mark task as complete
                self.task_manager.mark_complete(task)
        
        # Combine results
        combined_result = self._combine_results(query, results, enhanced_context, session_id)
        
        # Store response in memory
        self.memory.store_output(session_id, combined_result)
        
        # Store in Titans memory if enabled
        if self.titans_memory_enabled and self.titans_adapter:
            try:
                self.titans_adapter.store_response(query, combined_result, session_id)
            except Exception as e:
                print(f"⚠️ Error storing response in Titans memory: {e}")
        
        return combined_result
        
    def _combine_results(
        self,
        query: str,
        results: List[Dict[str, Any]],
        enhanced_context: str,
        session_id: str
    ) -> str:
        """
        Combine task results into a coherent response.
        
        Args:
            query: Original user query
            results: List of task results
            enhanced_context: Enhanced context from Titans memory
            session_id: Session identifier
            
        Returns:
            Combined response
        """
        # Get conversation context
        conversation_context = self.memory.get_context(session_id)
        
        # Create prompt with context and results
        results_text = ""
        for i, result in enumerate(results, 1):
            task = result.get("task", "Unknown task")
            tool = result.get("tool", "Unknown tool")
            
            if "result" in result:
                # Truncate long results
                result_text = result["result"]
                if len(result_text) > 1000:
                    result_text = result_text[:1000] + "... [truncated]"
                    
                results_text += f"Result {i}:\n"
                results_text += f"Task: {task}\n"
                results_text += f"Tool: {tool}\n"
                results_text += f"Output: {result_text}\n\n"
            elif "error" in result:
                results_text += f"Result {i} (Error):\n"
                results_text += f"Task: {task}\n"
                results_text += f"Tool: {tool}\n"
                results_text += f"Error: {result['error']}\n\n"
        
        prompt = [
            {"role": "system", "content": f"""You are {self.agent_name}, an advanced AI assistant.
Your task is to synthesize the results of multiple subtasks into a coherent response to the user's original query.

Original query: {query}

{enhanced_context}

Conversation history:
{conversation_context}

Task results:
{results_text}

Create a comprehensive, well-structured response that addresses the user's original query.
Focus on providing a clear, concise, and helpful answer based on the task results.
Do not mention the individual tasks or tools unless relevant to the user's understanding.
Do not include phrases like "Based on the results" or "According to the task results".
Write as if you are directly answering the user's query with your own knowledge."""},
            {"role": "user", "content": query}
        ]
        
        # Get response from LLM
        response = self.llm.chat(prompt)
        
        return response
        
    def get_agent_identity(self) -> Dict[str, Any]:
        """
        Get agent identity information.
        
        Returns:
            Dictionary with agent identity information
        """
        return {
            "name": self.agent_name,
            "version": self.agent_version,
            "creation_date": self.agent_creation_date,
            "capabilities": [
                "Question answering",
                "Task planning and execution",
                "Code generation",
                "Content creation",
                "Web search",
                "Web scraping"
            ],
            "tools": list(self.tools.keys()),
            "memory_enabled": True,
            "titans_memory_enabled": self.titans_memory_enabled
        }
        
    def get_memory_insights(self) -> Dict[str, Any]:
        """
        Get insights about the agent's memory.
        
        Returns:
            Dictionary with memory insights
        """
        insights = {
            "basic_memory": {
                "type": "SQLite",
                "path": self.memory_path
            }
        }
        
        if self.titans_memory_enabled and self.titans_adapter:
            try:
                titans_insights = self.titans_adapter.get_insights()
                insights["adaptive_memory"] = {
                    "type": "Titans",
                    "enabled": True,
                    "insights": titans_insights
                }
            except Exception as e:
                insights["adaptive_memory"] = {
                    "type": "Titans",
                    "enabled": True,
                    "error": str(e)
                }
        else:
            insights["adaptive_memory"] = {
                "type": "Titans",
                "enabled": False
            }
            
        return insights