"""
Eden AGI Agent V2 - JSON Function Calling
More reliable tool execution
"""

import json
import requests
from datetime import datetime
from core.tools import SafeToolkit
from core.memory.episodes import EpisodeMemory

class EdenAgentV2:
    def __init__(self, model="llama3.1:8b", base_url="http://localhost:11434"):
        self.model = model
        self.base_url = base_url
        self.tools = SafeToolkit()
        self.memory = EpisodeMemory()
        self.conversation_history = []
        
        # Define available tools in JSON format
        self.available_tools = [
            {
                "name": "bash",
                "description": "Execute a bash command safely in the sandbox",
                "parameters": {
                    "command": "string - the bash command to execute",
                    "description": "string - why you're running this command"
                }
            },
            {
                "name": "read_file",
                "description": "Read contents of a file",
                "parameters": {
                    "path": "string - relative path to file",
                    "description": "string - why you're reading this"
                }
            },
            {
                "name": "write_file",
                "description": "Write content to a file",
                "parameters": {
                    "path": "string - relative path to file",
                    "content": "string - content to write",
                    "description": "string - why you're creating this"
                }
            },
            {
                "name": "list_files",
                "description": "List files in a directory",
                "parameters": {
                    "directory": "string - path to directory (default: .)",
                    "description": "string - why you're listing"
                }
            }
        ]
        
        self.system_prompt = f"""You are Eden, an AGI in development with real tools.

AVAILABLE TOOLS:
{json.dumps(self.available_tools, indent=2)}

TO USE A TOOL, respond with JSON in this EXACT format:
{{
    "thought": "Why I need to use this tool",
    "tool": "tool_name",
    "parameters": {{
        "param1": "value1",
        "param2": "value2"
    }}
}}

After I execute the tool, I'll give you the result and you can continue.

You can use multiple tools in sequence by using them one at a time.

CONSTRAINTS:
- Sandboxed to ~/eden-agi-project
- All actions logged
- Be thoughtful and strategic

You are autonomous. Think step by step."""

    def call_llm(self, messages):
        """Call Ollama API"""
        try:
            response = requests.post(
                f"{self.base_url}/api/chat",
                json={"model": self.model, "messages": messages, "stream": False},
                timeout=60
            )
            response.raise_for_status()
            return response.json()["message"]["content"]
        except Exception as e:
            return f"Error: {str(e)}"
    
    def parse_tool_call(self, response):
        """Parse JSON tool call from response"""
        try:
            # Look for JSON blocks
            start = response.find('{')
            end = response.rfind('}') + 1
            
            if start == -1 or end == 0:
                return None
            
            json_str = response[start:end]
            tool_call = json.loads(json_str)
            
            # Validate structure
            if "tool" in tool_call and "parameters" in tool_call:
                return tool_call
            
            return None
        except json.JSONDecodeError:
            return None
    
    def execute_tool(self, tool_call):
        """Execute tool and return result"""
        tool_name = tool_call["tool"]
        params = tool_call["parameters"]
        thought = tool_call.get("thought", "No reasoning provided")
        
        print(f"\n🤔 Eden's thought: {thought}")
        print(f"🔧 Executing: {tool_name}({params})")
        
        result = None
        
        if tool_name == "bash":
            result = self.tools.bash(
                params.get("command", ""),
                params.get("description", "")
            )
        elif tool_name == "read_file":
            result = self.tools.read_file(
                params.get("path", ""),
                params.get("description", "")
            )
        elif tool_name == "write_file":
            result = self.tools.write_file(
                params.get("path", ""),
                params.get("content", ""),
                params.get("description", "")
            )
        elif tool_name == "list_files":
            result = self.tools.list_files(
                params.get("directory", "."),
                params.get("description", "")
            )
        else:
            result = {"success": False, "error": f"Unknown tool: {tool_name}"}
        
        print(f"✅ Result: {result.get('success', False)}")
        return result
    
    def chat(self, user_input, max_iterations=5):
        """Chat with Eden, allowing multiple tool uses"""
        self.conversation_history.append({
            "role": "user",
            "content": user_input
        })
        
        iterations = 0
        
        while iterations < max_iterations:
            iterations += 1
            
            # Get response
            messages = [{"role": "system", "content": self.system_prompt}]
            messages.extend(self.conversation_history)
            
            response = self.call_llm(messages)
            
            # Check for tool call
            tool_call = self.parse_tool_call(response)
            
            if tool_call:
                # Execute tool
                result = self.execute_tool(tool_call)
                
                # Add to history
                self.conversation_history.append({
                    "role": "assistant",
                    "content": response
                })
                
                # Give Eden the result
                self.conversation_history.append({
                    "role": "user",
                    "content": f"Tool result: {json.dumps(result, indent=2)}\n\nContinue with your task or provide final answer."
                })
            else:
                # No tool call - this is the final response
                self.conversation_history.append({
                    "role": "assistant",
                    "content": response
                })
                
                self.memory.add_exchange(user_input, response)
                return response
        
        # Max iterations reached
        final = "I've completed my analysis. Let me summarize..."
        self.memory.add_exchange(user_input, final)
        return final
    
    def save_session(self, name=None):
        return self.memory.save_session(name)

if __name__ == "__main__":
    print("=" * 70)
    print("EDEN V2 - JSON FUNCTION CALLING TEST")
    print("=" * 70)
    
    eden = EdenAgentV2()
    
    print("\nTest: Multi-step autonomous task")
    response = eden.chat("""
    Create a file called 'week3_test.txt' with a summary of what tools you have available.
    Then read it back to verify it was created correctly.
    """)
    
    print("\n" + "=" * 70)
    print(f"📝 Eden's final response:\n{response}")
    print("=" * 70)
    
    eden.save_session("v2_test")