"""
OpenRouter AI Agent with Tool Calling Support
Supports the OpenRouter API protocol for function/tool calling
"""

import json
import requests
import time
from typing import Dict, List, Any, Optional, Callable, Union
from dataclasses import dataclass, field
from enum import Enum
import inspect
from functools import wraps
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class Role(Enum):
    """Message roles for the conversation"""
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"


@dataclass
class Message:
    """Message structure for the conversation"""
    role: Role
    content: Optional[str] = None
    name: Optional[str] = None
    tool_calls: Optional[List[Dict[str, Any]]] = None
    tool_call_id: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert message to OpenRouter API format"""
        message_dict = {
            "role": self.role.value,
            "content": self.content if self.content is not None else ""
        }
        
        if self.name:
            message_dict["name"] = self.name
            
        if self.role == Role.ASSISTANT and self.tool_calls:
            message_dict["tool_calls"] = self.tool_calls
            
        if self.role == Role.TOOL and self.tool_call_id:
            message_dict["tool_call_id"] = self.tool_call_id
            
        return message_dict


@dataclass
class Tool:
    """Tool definition for function calling"""
    name: str
    description: str
    parameters: Dict[str, Any]
    function: Callable
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert tool to OpenRouter API format"""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters
            }
        }


class OpenRouterAgent:
    """AI Agent using OpenRouter API with tool calling capability"""
    
    def __init__(
        self,
        api_key: str,
        model: str = "qwen/qwen3-coder-next",
        base_url: str = "https://openrouter.ai/api/v1",
        temperature: float = 0.7,
        max_tokens: int = 2048,
        max_iterations: int = 10
    ):
        """
        Initialize the OpenRouter agent.
        
        Args:
            api_key: OpenRouter API key
            model: Model identifier (default: "openai/gpt-4-turbo")
            base_url: OpenRouter API base URL
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            max_iterations: Maximum tool-calling iterations
        """
        self.api_key = api_key
        self.model = model
        self.base_url = base_url
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.max_iterations = max_iterations
        
        self.tools: Dict[str, Tool] = {}
        self.conversation_history: List[Message] = []
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://github.com",
            "X-Title": "AI Agent"
        }
        
        self._register_builtin_tools()
    
    def register_tool(
        self,
        name: str,
        description: str,
        parameters: Dict[str, Any]
    ) -> Callable:
        """Decorator to register a tool function"""
        def decorator(func: Callable) -> Callable:
            tool = Tool(
                name=name,
                description=description,
                parameters=parameters,
                function=func
            )
            self.tools[name] = tool
            return func
        return decorator
    
    def _make_api_call(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """Make API call to OpenRouter"""
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens
        }
        
        if tools:
            payload["tools"] = tools
        
        response = requests.post(
            f"{self.base_url}/chat/completions",
            headers=self.headers,
            json=payload,
            timeout=30
        )
        response.raise_for_status()
        return response.json()
    
    def _process_tool_calls(self, tool_calls: List[Dict]) -> List[Message]:
        """Process tool calls and return tool result messages"""
        tool_messages = []
        
        for tool_call in tool_calls:
            func_name = tool_call["function"]["name"]
            func_args = json.loads(tool_call["function"]["arguments"])
            
            if func_name in self.tools:
                func = self.tools[func_name].function
                try:
                    result = func(**func_args)
                except Exception as e:
                    result = f"Error: {str(e)}"
            else:
                result = f"Error: Unknown tool {func_name}"
            
            tool_messages.append(
                Message(
                    role=Role.TOOL,
                    content=str(result),
                    tool_call_id=tool_call["id"]
                )
            )
        
        return tool_messages
        
    def _register_builtin_tools(self) -> None:
        """Register built-in tools for the agent"""
        
        @self.register_tool(
            name="get_current_time",
            description="Get the current date and time",
            parameters={
                "type": "object",
                "properties": {},
                "required": []
            }
        )
        def get_current_time() -> str:
            """Returns current date and time as ISO format string"""
            from datetime import datetime
            return datetime.now().isoformat()
        
        @self.register_tool(
            name="calculate",
            description="Perform mathematical calculations",
            parameters={
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "Mathematical expression to evaluate"
                    }
                },
                "required": ["expression"]
            }
        )
        def calculate(expression: str) -> str:
            """Evaluate a mathematical expression safely"""
            try:
                allowed_chars = set("0123456789+-*/.() ")
                if not all(c in allowed_chars for c in expression):
                    return "Error: Expression contains disallowed characters"
                
                result = eval(expression, {"__builtins__": {}}, {})
                return f"{expression} = {result}"
            except Exception as e:
                return f"Error evaluating expression: {str(e)}"
        
        @self.register_tool(
            name="read_file",
            description="Read the contents of a file",
            parameters={
                "type": "object",
                "properties": {
                    "file_path": {
                        "type": "string",
                        "description": "The absolute path to the file to read"
                    }
                },
                "required": ["file_path"]
            }
        )
        def read_file(file_path: str) -> str:
            """Read and return the contents of a file"""
            try:
                with open(file_path, "r", encoding="utf-8") as file:
                    return file.read()
            except FileNotFoundError:
                return f"Error: File not found: {file_path}"
            except PermissionError:
                return f"Error: Permission denied: {file_path}"
            except Exception as e:
                return f"Error reading file {file_path}: {str(e)}"
        
        @self.register_tool(
            name="write_file",
            description="Write content to a file",
            parameters={
                "type": "object",
                "properties": {
                    "file_path": {
                        "type": "string",
                        "description": "The absolute path to the file to write"
                    },
                    "content": {
                        "type": "string",
                        "description": "The content to write to the file"
                    }
                },
                "required": ["file_path", "content"]
            }
        )
        def write_file(file_path: str, content: str) -> str:
            """Write content to a file"""
            try:
                with open(file_path, "w", encoding="utf-8") as file:
                    file.write(content)
                return f"Successfully wrote to {file_path}"
            except PermissionError:
                return f"Error: Permission denied: {file_path}"
            except Exception as e:
                return f"Error writing to file {file_path}: {str(e)}"
        
        @self.register_tool(
            name="list_directory",
            description="List files and directories in a given path",
            parameters={
                "type": "object",
                "properties": {
                    "directory_path": {
                        "type": "string",
                        "description": "The absolute path to the directory to list"
                    }
                },
                "required": ["directory_path"]
            }
        )
        def list_directory(directory_path: str) -> str:
            """List files and directories in a path"""
            import os
            try:
                entries = os.listdir(directory_path)
                if not entries:
                    return f"Directory is empty: {directory_path}"
                result = []
                for entry in entries:
                    full_path = os.path.join(directory_path, entry)
                    if os.path.isdir(full_path):
                        result.append(f"{entry}/")
                    else:
                        result.append(entry)
                return "\n".join(sorted(result))
            except FileNotFoundError:
                return f"Error: Directory not found: {directory_path}"
            except PermissionError:
                return f"Error: Permission denied: {directory_path}"
            except Exception as e:
                return f"Error listing directory {directory_path}: {str(e)}"
        
        @self.register_tool(
            name="file_exists",
            description="Check if a file or directory exists",
            parameters={
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "The absolute path to check"
                    }
                },
                "required": ["path"]
            }
        )
        def file_exists(path: str) -> str:
            """Check if a file or directory exists"""
            import os
            if os.path.exists(path):
                if os.path.isfile(path):
                    return f"File exists: {path}"
                elif os.path.isdir(path):
                    return f"Directory exists: {path}"
            return f"Path does not exist: {path}"
        
        @self.register_tool(
            name="ask_user",
            description="Ask the user a question and get their response",
            parameters={
                "type": "object",
                "properties": {
                    "question": {
                        "type": "string",
                        "description": "The question to ask the user"
                    }
                },
                "required": ["question"]
            }
        )
        def ask_user(question: str) -> str:
            """Ask the user a question and return their response"""
            return f"[USER_INPUT_NEEDED]: {question}"
    
    def read_file(self, file_path: str) -> str:
        """
        Read a file and return its content as a string.
        
        Args:
            file_path: The absolute path to the file to read
            
        Returns:
            The file content as a string or an error message
        """
        try:
            # Verify the file exists
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            return content
        except Exception as e:
            return f"Error reading file {file_path}: {str(e)}"
    
    def add_message(self, message: Message) -> None:
        """Add a message to the conversation history"""
        self.conversation_history.append(message)
    
    def run(
        self,
        user_input: str,
        system_prompt: Optional[str] = None,
        reset_conversation: bool = False
    ) -> str:
        """
        Main agentic loop: Process user input and handle tool calls
        
        Args:
            user_input: User's message
            system_prompt: Optional system prompt
            reset_conversation: Whether to reset conversation history
            
        Returns:
            Final assistant response
        """
        if reset_conversation:
            self.conversation_history = []
        
        # Add system prompt if provided
        if system_prompt and not any(m.role == Role.SYSTEM for m in self.conversation_history):
            self.add_message(Message(role=Role.SYSTEM, content=system_prompt))
        
        # Add user message
        self.add_message(Message(role=Role.USER, content=user_input))
        
        iteration = 0
        final_response = None
        
        while iteration < self.max_iterations:
            iteration += 1
            logger.info(f"Iteration {iteration}/{self.max_iterations}")
            
            # Prepare messages for API call
            messages = [msg.to_dict() for msg in self.conversation_history]
            
            # Prepare tools list if we have any
            tools_list = None
            if self.tools:
                tools_list = [tool.to_dict() for tool in self.tools.values()]
            
            # Make API call
            response = self._make_api_call(messages, tools_list)
            
            # Extract the response
            choice = response["choices"][0]
            message = choice["message"]
            
            # Check if there are tool calls
            if "tool_calls" in message and message["tool_calls"]:
                # Create assistant message with tool calls
                assistant_message = Message(
                    role=Role.ASSISTANT,
                    content=message.get("content"),
                    tool_calls=message["tool_calls"]
                )
                
                self.add_message(assistant_message)
                
                # Process tool calls and get results
                tool_messages = self._process_tool_calls(message["tool_calls"])
                
                # Add tool results to conversation
                for tool_msg in tool_messages:
                    self.add_message(tool_msg)
                
                # Continue the loop to send tool results back to the model
                continue
            else:
                # No tool calls, this is the final response
                assistant_message = Message(
                    role=Role.ASSISTANT,
                    content=message.get("content", "")
                )
                
                self.add_message(assistant_message)
                final_response = message.get("content", "")
                break
        
        if final_response is None:
            final_response = "Maximum iterations reached without final response."
        
        return final_response
    
    def get_conversation_history(self) -> List[Dict[str, Any]]:
        """Get the conversation history in a readable format"""
        return [msg.to_dict() for msg in self.conversation_history]
    
    def clear_history(self) -> None:
        """Clear the conversation history"""
        self.conversation_history = []


# Example custom tools
def create_custom_tools_agent(api_key: str) -> OpenRouterAgent:
    """Create an agent with custom tools for demonstration"""
    
    agent = OpenRouterAgent(
        api_key=api_key,
        model="qwen/qwen3-coder-next",  # You can change to any OpenRouter model
        temperature=0.7,
        max_iterations=8
    )
    
    # Register custom tools using decorator
    @agent.register_tool(
        name="get_weather",
        description="Get weather information for a city",
        parameters={
            "type": "object",
            "properties": {
                "city": {
                    "type": "string",
                    "description": "The city name"
                },
                "country": {
                    "type": "string",
                    "description": "The country code (optional)",
                    "default": "US"
                }
            },
            "required": ["city"]
        }
    )
    def get_weather(city: str, country: str = "US") -> str:
        """Simulated weather function - in production, connect to real API"""
        # In a real implementation, you would call a weather API here
        weather_data = {
            "New York": {"temp": 72, "condition": "Sunny", "humidity": 65},
            "London": {"temp": 60, "condition": "Cloudy", "humidity": 80},
            "Tokyo": {"temp": 75, "condition": "Rainy", "humidity": 90},
            "Sydney": {"temp": 85, "condition": "Clear", "humidity": 70},
        }
        
        if city in weather_data:
            data = weather_data[city]
            return f"Weather in {city}, {country}: {data['temp']}°F, {data['condition']}, Humidity: {data['humidity']}%"
        else:
            return f"Weather data not available for {city}. Try: New York, London, Tokyo, or Sydney."
    
    @agent.register_tool(
        name="search_web",
        description="Search the web for information",
        parameters={
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query"
                },
                "num_results": {
                    "type": "integer",
                    "description": "Number of results to return",
                    "default": 3
                }
            },
            "required": ["query"]
        }
    )
    def search_web(query: str, num_results: int = 3) -> str:
        """Simulated web search - in production, connect to search API"""
        # Simulated search results
        results = [
            f"1. {query} - Recent news article",
            f"2. {query} - Wikipedia entry",
            f"3. {query} - Research paper",
            f"4. {query} - Blog post",
            f"5. {query} - Official documentation"
        ]
        
        return f"Search results for '{query}':\n" + "\n".join(results[:num_results])
    
    @agent.register_tool(
        name="format_json",
        description="Format data as JSON string",
        parameters={
            "type": "object",
            "properties": {
                "data": {
                    "type": "string",
                    "description": "Data to format as JSON"
                }
            },
            "required": ["data"]
        }
    )
    def format_json(data: str) -> str:
        """Format data as JSON string"""
        try:
            # Try to parse and pretty-print if it's JSON-like
            parsed = json.loads(data)
            return json.dumps(parsed, indent=2)
        except:
            # If not JSON, just return as-is
            return data
    
    return agent


def main():
    """Example usage of the OpenRouter agent"""
    
    # Replace with your actual OpenRouter API key
    API_KEY = ""
    
    if API_KEY == "your-openrouter-api-key-here":
        print("Please set your OpenRouter API key in the code")
        return
    
    # Create agent with custom tools
    agent = create_custom_tools_agent(API_KEY)
    
    # Example system prompt
    system_prompt = """You are a helpful AI assistant with access to various tools.
    When asked a question, use the available tools to gather information or perform actions.
    Always be concise and helpful in your responses."""
    
    print("OpenRouter AI Agent with Tool Calling")
    print("=" * 50)
    print("Type 'quit' to exit, 'clear' to clear history, 'history' to view conversation")
    print()
    
    while True:
        try:
            user_input = input("\nYou: ").strip()
            
            if user_input.lower() == 'quit':
                print("Goodbye!")
                break
            elif user_input.lower() == 'clear':
                agent.clear_history()
                print("Conversation history cleared.")
                continue
            elif user_input.lower() == 'history':
                history = agent.get_conversation_history()
                print("\nConversation History:")
                for i, msg in enumerate(history):
                    role = msg.get('role', 'unknown').upper()
                    content = msg.get('content', '')[:100] + "..." if len(msg.get('content', '')) > 100 else msg.get('content', '')
                    print(f"{i+1}. [{role}]: {content}")
                continue
            
            if not user_input:
                continue
            
            print("\nAssistant: ", end="", flush=True)
            
            # Run the agent
            response = agent.run(
                user_input=user_input,
                system_prompt=system_prompt,
                reset_conversation=False  # Set to True to start fresh each time
            )
            
            print(response)
            
        except KeyboardInterrupt:
            print("\n\nInterrupted. Exiting...")
            break
        except Exception as e:
            print(f"\nError: {e}")
            logger.exception("An error occurred")


if __name__ == "__main__":
    main()