"""
Core AI Agent Implementation

This module contains the main Agent class that orchestrates all operations,
including reasoning, tool selection, memory management, and output formatting.
"""

import json
import time
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
from enum import Enum
import openai
from openai import OpenAI

from memory.memory_manager import MemoryManager
from tools.tool_registry import ToolRegistry
from utils.formatter import OutputFormatter


class ReasoningStep(Enum):
    """Enumeration of reasoning steps for transparency"""
    TASK_ANALYSIS = "task_analysis"
    TOOL_SELECTION = "tool_selection"
    EXECUTION = "execution"
    SYNTHESIS = "synthesis"
    EVALUATION = "evaluation"


@dataclass
class AgentResponse:
    """Structured response from the agent"""
    reasoning_steps: List[Dict[str, Any]]
    tool_used: Optional[str]
    tool_output: Any
    final_response: str
    execution_time: float
    memory_updates: List[str]


class AIAgent:
    """
    Main AI Agent class implementing the orchestrator-worker pattern
    with memory, tool integration, and visible reasoning.
    """
    
    def __init__(self, 
                 model_name: str = "gpt-4",
                 memory_persist_path: str = "./agent_memory",
                 enable_reasoning_display: bool = True,
                 realtime_output: bool = True,
                 minimal_output: bool = False,
                 typing_delay: float = 0.01):
        """
        Initialize the AI Agent
        
        Args:
            model_name: OpenAI model to use
            memory_persist_path: Path to persist memory data
            enable_reasoning_display: Whether to show detailed reasoning
            realtime_output: Whether to print output in real-time as steps are processed
            minimal_output: Whether to show minimal output (less info)
            typing_delay: Delay between characters for typewriter effect
        """
        self.client = OpenAI()
        self.model_name = model_name
        self.enable_reasoning_display = enable_reasoning_display
        self.realtime_output = realtime_output
        self.minimal_output = minimal_output
        self.typing_delay = typing_delay
        
        # Initialize components
        self.memory_manager = MemoryManager(persist_path=memory_persist_path)
        self.tool_registry = ToolRegistry()
        self.formatter = OutputFormatter(minimal=minimal_output, typing_delay=typing_delay)
        
        # Reasoning state
        self.current_reasoning_steps = []
        
        # System prompt for the agent
        self.system_prompt = """
You are an advanced AI agent with access to memory and various tools. Your role is to:

1. Analyze user requests carefully
2. Break down complex tasks into manageable steps
3. Select and use appropriate tools
4. Maintain context through memory
5. Provide clear, reasoned responses

Available capabilities:
- Long-term memory for storing and retrieving information
- Web browsing and content extraction
- Python code execution
- File system operations
- Screenshot capture
- Content generation and analysis
- Diff checking and comparison

Always explain your reasoning process step by step, show which tools you're using, and provide formatted output.
"""

    def process_query(self, user_query: str, context: Optional[Dict] = None) -> AgentResponse:
        """
        Main method to process user queries with full reasoning display
        
        Args:
            user_query: The user's input query
            context: Optional additional context
            
        Returns:
            AgentResponse with detailed reasoning and results
        """
        start_time = time.time()
        self.current_reasoning_steps = []
        
        try:
            # Step 1: Task Analysis
            self._add_reasoning_step(ReasoningStep.TASK_ANALYSIS, 
                                   "Analyzing user query and determining approach")
            
            task_analysis = self._analyze_task(user_query, context)
            
            # Step 2: Memory Retrieval
            relevant_memory = self.memory_manager.retrieve_relevant_context(user_query)
            if relevant_memory:
                self._add_reasoning_step("memory_retrieval", 
                                       f"Retrieved {len(relevant_memory)} relevant memory items")
            
            # Step 3: Tool Selection
            self._add_reasoning_step(ReasoningStep.TOOL_SELECTION,
                                   "Selecting appropriate tools for the task")
            
            selected_tools = self._select_tools(task_analysis, user_query)
            
            # Step 4: Tool Execution
            tool_outputs = {}
            for tool_name in selected_tools:
                self._add_reasoning_step(ReasoningStep.EXECUTION,
                                       f"Executing tool: {tool_name}")
                
                tool_output = self._execute_tool(tool_name, user_query, task_analysis)
                tool_outputs[tool_name] = tool_output
                if self.realtime_output:
                    color = 'green' if (isinstance(tool_output, dict) and tool_output.get('success')) else 'red'
                    self.formatter.section_start(f"🔧 Tool Output: {tool_name}", color=color)
                    if not self.minimal_output:
                        self.formatter.typewriter_print(self.formatter._format_tool_output(tool_output))
                    self.formatter.section_end(keep_title=True)
            
            # Step 5: Synthesis
            self._add_reasoning_step(ReasoningStep.SYNTHESIS,
                                   "Synthesizing results and generating response")
            
            final_response = self._synthesize_response(
                user_query, task_analysis, tool_outputs, relevant_memory
            )
            if self.realtime_output:
                self.formatter.section_start("💬 Final Response", color='header')
                self.formatter.typewriter_print(self.formatter._format_text_content(final_response))
                self.formatter.section_end(keep_title=True)
            # --- Enhancement: Execute generated code if present in final_response ---
            executed_code_result = None
            if 'python_tool' in selected_tools:
                import re
                # Extract all code blocks (with or without python tag)
                code_blocks = re.findall(r'```(?:python)?[ \n]*(.*?)[\n`]*```', final_response, re.DOTALL)
                cleaned_code = None
                if code_blocks:
                    # Prefer code blocks with function definitions or assignments
                    for block in code_blocks:
                        block_clean = block.strip()
                        if 'def ' in block_clean or 'return ' in block_clean or '=' in block_clean:
                            cleaned_code = block_clean
                            break
                    if not cleaned_code:
                        cleaned_code = code_blocks[0].strip()
                # Fallback: try to extract inline code
                if not cleaned_code:
                    inline_code = re.findall(r'`([^`]+)`', final_response)
                    for code in inline_code:
                        if 'def ' in code or 'return ' in code or '=' in code:
                            cleaned_code = code.strip()
                            break
                # Only execute and add to tool_outputs if we actually found code
                if cleaned_code:
                    python_tool = self.tool_registry.get_tool('python_tool')
                    if python_tool:
                        executed_code_result = python_tool.execute(cleaned_code, task_analysis)
                        tool_outputs['python_tool_execution'] = executed_code_result
                        self._add_reasoning_step('code_execution', 'Executed generated Python code and captured result.')
                        if self.realtime_output:
                            color = 'green' if (isinstance(executed_code_result, dict) and executed_code_result.get('success')) else 'red'
                            self.formatter.section_start("🔧 Tool Output: python_tool_execution", color=color)
                            if not self.minimal_output:
                                self.formatter.typewriter_print(self.formatter._format_tool_output(executed_code_result))
                            self.formatter.section_end(keep_title=True)
            # --- End enhancement ---
            
            # Step 6: Memory Update
            memory_updates = self._update_memory(user_query, final_response, tool_outputs)
            
            # Step 7: Evaluation
            self._add_reasoning_step(ReasoningStep.EVALUATION,
                                   "Evaluating response quality and completeness")
            
            execution_time = time.time() - start_time
            
            return AgentResponse(
                reasoning_steps=self.current_reasoning_steps,
                tool_used=selected_tools[0] if selected_tools else None,
                tool_output=tool_outputs,
                final_response=final_response,
                execution_time=execution_time,
                memory_updates=memory_updates
            )
            
        except Exception as e:
            self._add_reasoning_step("error", f"Error occurred: {str(e)}")
            return AgentResponse(
                reasoning_steps=self.current_reasoning_steps,
                tool_used=None,
                tool_output=None,
                final_response=f"I encountered an error while processing your request: {str(e)}",
                execution_time=time.time() - start_time,
                memory_updates=[]
            )

    def _add_reasoning_step(self, step_type: str, description: str, details: Optional[Dict] = None):
        """Add a reasoning step to the log and print if realtime_output is enabled"""
        step = {
            'step_type': step_type,
            'description': description,
            'timestamp': time.time(),
        }
        if details:
            step['details'] = details
        self.current_reasoning_steps.append(step)
        if getattr(self, 'realtime_output', False):
            color = 'cyan' if 'analysis' in str(step_type).lower() else (
                'yellow' if 'tool' in str(step_type).lower() else (
                'green' if 'synthesis' in str(step_type).lower() else (
                'red' if 'error' in str(step_type).lower() else 'blue')))
            self.formatter.section_start(f"🧠 {step_type}", color=color)
            if not self.minimal_output:
                self.formatter.typewriter_print(self.formatter._format_reasoning_step(len(self.current_reasoning_steps), step))
            self.formatter.section_end(keep_title=True)

    def _analyze_task(self, user_query: str, context: Optional[Dict] = None) -> Dict[str, Any]:
        """Analyze the user's task to determine approach"""
        
        analysis_prompt = f"""
Analyze this user query and provide a structured analysis:

Query: {user_query}
Context: {context or 'None'}

Provide analysis in this JSON format:
{{
    "task_type": "information_retrieval|code_execution|content_generation|web_interaction|file_operation|analysis",
    "complexity": "simple|moderate|complex",
    "required_tools": ["tool1", "tool2"],
    "subtasks": ["subtask1", "subtask2"],
    "expected_output_format": "text|code|image|document|data"
}}
"""
        
        try:
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": self.system_prompt},
                    {"role": "user", "content": analysis_prompt}
                ],
                temperature=0.1
            )
            
            analysis_text = response.choices[0].message.content
            # Try to extract JSON from the response
            import re
            json_match = re.search(r'\{.*\}', analysis_text, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
            else:
                # Fallback analysis
                return {
                    "task_type": "general",
                    "complexity": "moderate",
                    "required_tools": [],
                    "subtasks": [user_query],
                    "expected_output_format": "text"
                }
                
        except Exception as e:
            self._add_reasoning_step("analysis_error", f"Task analysis failed: {str(e)}")
            return {
                "task_type": "general",
                "complexity": "moderate", 
                "required_tools": [],
                "subtasks": [user_query],
                "expected_output_format": "text"
            }

    def _select_tools(self, task_analysis: Dict[str, Any], user_query: str) -> List[str]:
        """Select appropriate tools based on task analysis"""
        
        # Get available tools
        available_tools = self.tool_registry.get_available_tools()
        
        # Simple tool selection logic based on keywords and task analysis
        selected_tools = []
        selection_reasons = []
        
        query_lower = user_query.lower()
        
        # Tool selection rules
        if any(word in query_lower for word in ['screenshot', 'capture', 'image']):
            if 'screenshot_tool' in available_tools:
                selected_tools.append('screenshot_tool')
                selection_reasons.append("Selected screenshot_tool because the query contains keywords related to screenshots or images.")
        
        if any(word in query_lower for word in ['python', 'code', 'execute', 'run', 'calculate']):
            if 'python_tool' in available_tools:
                selected_tools.append('python_tool')
                selection_reasons.append("Selected python_tool because the query contains keywords related to code execution or calculation.")
        
        if any(word in query_lower for word in ['web', 'browse', 'website', 'url', 'internet']):
            if 'web_browser_tool' in available_tools:
                selected_tools.append('web_browser_tool')
                selection_reasons.append("Selected web_browser_tool because the query contains keywords related to web browsing or URLs.")
        
        if any(word in query_lower for word in ['file', 'read', 'write', 'codebase']):
            if 'file_system_tool' in available_tools:
                selected_tools.append('file_system_tool')
                selection_reasons.append("Selected file_system_tool because the query contains keywords related to files or codebase operations.")
        
        if any(word in query_lower for word in ['scrape', 'crawl', 'extract']):
            if 'web_scraper_tool' in available_tools:
                selected_tools.append('web_scraper_tool')
                selection_reasons.append("Selected web_scraper_tool because the query contains keywords related to scraping or extraction.")
        
        if any(word in query_lower for word in ['diff', 'compare', 'difference']):
            if 'diff_checker_tool' in available_tools:
                selected_tools.append('diff_checker_tool')
                selection_reasons.append("Selected diff_checker_tool because the query contains keywords related to comparison or differences.")
        
        if any(word in query_lower for word in ['write', 'book', 'document', 'article']):
            if 'content_generator_tool' in available_tools:
                selected_tools.append('content_generator_tool')
                selection_reasons.append("Selected content_generator_tool because the query contains keywords related to content generation or writing.")
        
        # Use task analysis suggestions
        for tool in task_analysis.get('required_tools', []):
            if tool in available_tools and tool not in selected_tools:
                selected_tools.append(tool)
                selection_reasons.append(f"Selected {tool} because it was suggested by the task analysis.")
        
        # Add reasoning step for tool selection
        if selection_reasons:
            self._add_reasoning_step("tool_selection_reasoning", "\n".join(selection_reasons))
        
        return selected_tools

    def _execute_tool(self, tool_name: str, user_query: str, task_analysis: Dict[str, Any]) -> Any:
        """Execute a specific tool"""
        
        tool = self.tool_registry.get_tool(tool_name)
        if not tool:
            return f"Tool {tool_name} not found"
        
        try:
            return tool.execute(user_query, task_analysis)
        except Exception as e:
            return f"Error executing {tool_name}: {str(e)}"

    def _synthesize_response(self, user_query: str, task_analysis: Dict[str, Any], 
                           tool_outputs: Dict[str, Any], relevant_memory: List[Dict]) -> str:
        """Synthesize final response from all available information"""
        
        synthesis_prompt = f"""
Based on the following information, provide a comprehensive response to the user's query:

User Query: {user_query}

Task Analysis: {json.dumps(task_analysis, indent=2)}

Tool Outputs: {json.dumps(tool_outputs, indent=2)}

Relevant Memory: {json.dumps(relevant_memory, indent=2)}

Provide a clear, helpful response that directly addresses the user's needs.
"""
        
        try:
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": self.system_prompt},
                    {"role": "user", "content": synthesis_prompt}
                ],
                temperature=0.3
            )
            
            return response.choices[0].message.content
            
        except Exception as e:
            return f"I processed your request but encountered an error during synthesis: {str(e)}"

    def _update_memory(self, user_query: str, response: str, tool_outputs: Dict[str, Any]) -> List[str]:
        """Update agent memory with new information"""
        
        updates = []
        
        # Store the interaction
        interaction_data = {
            "query": user_query,
            "response": response,
            "tool_outputs": tool_outputs,
            "timestamp": time.time()
        }
        
        memory_id = self.memory_manager.store_interaction(interaction_data)
        updates.append(f"Stored interaction with ID: {memory_id}")
        
        # Extract and store key information from tool outputs
        for tool_name, output in tool_outputs.items():
            if output and isinstance(output, str) and len(output) > 100:
                # Store significant tool outputs as separate memory items
                memory_id = self.memory_manager.store_knowledge(
                    content=output,
                    source=f"tool_{tool_name}",
                    metadata={"query": user_query, "tool": tool_name}
                )
                updates.append(f"Stored {tool_name} output with ID: {memory_id}")
        
        return updates

    def get_formatted_response(self, response: AgentResponse) -> str:
        """Get a formatted version of the agent response"""
        return self.formatter.format_agent_response(response)

    def clear_memory(self):
        """Clear all stored memory"""
        self.memory_manager.clear_all_memory()

    def get_memory_stats(self) -> Dict[str, Any]:
        """Get statistics about stored memory"""
        return self.memory_manager.get_stats()

