"""
OpenRouter Agent implementation.
"""

import json
import logging
import requests
from typing import Any, Callable, Dict, List, Optional

from core.models import Message, Role, Tool

logger = logging.getLogger(__name__)


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,
        tools: Optional[Dict[str, Tool]] = None
    ):
        """
        Initialize the OpenRouter agent.
        
        Args:
            api_key: OpenRouter API key
            model: Model identifier (default: "qwen/qwen3-coder-next")
            base_url: OpenRouter API base URL
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            max_iterations: Maximum tool-calling iterations
            tools: Optional dictionary of pre-registered tools
        """
        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] = tools if tools else {}
        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"
        }
    
    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 add_tool(self, tool: Tool) -> None:
        """Add a tool to the agent"""
        self.tools[tool.name] = tool
    
    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 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 = []
        
        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))
        
        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}")
            
            messages = [msg.to_dict() for msg in self.conversation_history]
            
            tools_list = None
            if self.tools:
                tools_list = [tool.to_dict() for tool in self.tools.values()]
            
            response = self._make_api_call(messages, tools_list)
            
            choice = response["choices"][0]
            message = choice["message"]
            
            if "tool_calls" in message and message["tool_calls"]:
                assistant_message = Message(
                    role=Role.ASSISTANT,
                    content=message.get("content"),
                    tool_calls=message["tool_calls"]
                )
                
                self.add_message(assistant_message)
                
                tool_messages = self._process_tool_calls(message["tool_calls"])
                
                for tool_msg in tool_messages:
                    self.add_message(tool_msg)
                
                continue
            else:
                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 = []