import requests
import urllib.parse
from typing import Dict, Any, Optional, List
import webbrowser

class KlavisAPI:
    """API Client for Klavis API."""
    
    def __init__(self, api_key: str, base_url: str = "https://api.klavis.ai"):
        self.api_key = api_key
        self.base_url = base_url
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
    
    def _make_request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
        """Make HTTP request with error handling."""
        url = f"{self.base_url}{endpoint}"
        response = requests.request(method, url, headers=self.headers, **kwargs)
        response.raise_for_status()
        return response.json()
    
    def create_mcp_instance(self, server_name: str, user_id: str, platform_name: str) -> Dict[str, str]:
        """Create MCP server instance."""
        data = {
            "serverName": server_name,
            "userId": user_id,
            "platformName": platform_name,
            "connectionType": "StreamableHttp"
        }
        result = self._make_request("POST", "/mcp-server/instance/create", json=data)
        print(f"✅ Created {server_name} MCP instance")
        return {
            'serverUrl': result['serverUrl'],
            'instanceId': result['instanceId']
        }
    
    def list_tools(self, server_url: str) -> Dict[str, Any]:
        """List all available tools for an MCP server."""
        params = {"connection_type": "StreamableHttp"}
        encoded_url = urllib.parse.quote(server_url, safe='')
        return self._make_request("GET", f"/mcp-server/list-tools/{encoded_url}", params=params)
    
    def _convert_mcp_tools_to_openai_format(self, mcp_tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Convert MCP tools format to OpenAI function calling format."""
        openai_tools = []
        
        for tool in mcp_tools:
            openai_tool = {
                "type": "function",
                "function": {
                    "name": tool.get("name", ""),
                    "description": tool.get("description", ""),
                    "parameters": tool.get("inputSchema", {})
                }
            }
            openai_tools.append(openai_tool)
        
        return openai_tools
    
    def call_tool(self, server_url: str, tool_name: str, 
                  tool_args: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Call tool on MCP server."""
        data = {
            "serverUrl": server_url,
            "toolName": tool_name,
            "toolArgs": tool_args or {},
            "connectionType": "StreamableHttp"
        }
        return self._make_request("POST", "/mcp-server/call-tool", json=data)
    
    def redirect_to_oauth(self, instance_id: str, server_name: str) -> None:
        """Open OAuth authorization URL in browser."""
        oauth_url = f"{self.base_url}/oauth/{server_name.lower()}/authorize?instance_id={instance_id}"
        print(f"🔐 Opening OAuth authorization for {server_name}, if you are not redirected, please open the following URL in your browser: {oauth_url}")
        webbrowser.open(oauth_url)

import json

class Agent:
    def __init__(self, fireworks_client, klavis_api_client, mcp_server_url):
        self.fireworks = fireworks_client
        self.klavis = klavis_api_client
        self.mcp_server_url = mcp_server_url
        self.model = "accounts/fireworks/models/qwen2p5-72b-instruct"
        print(f"🤖 Agent initialized with model: {self.model}")
    
    def process_request(self, user_message):
        """Process a user request using Fireworks AI + Klavis integration."""
        # 1. Get available tools
        mcp_tools = self.klavis.list_tools(self.mcp_server_url)
        tools = self.klavis._convert_mcp_tools_to_openai_format(mcp_tools.get('tools', []))
        
        # 2. Call LLM with tools
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": user_message}
        ]
        
        response = self.fireworks.chat.completions.create(
            model=self.model,
            messages=messages,
            tools=tools
        )
        
        assistant_message = response.choices[0].message
        messages.append(assistant_message)
        
        # 3. If LLM wants to use tools
        if assistant_message.tool_calls:
            
            # Execute each tool call
            for tool_call in assistant_message.tool_calls:
                tool_name = tool_call.function.name
                tool_args = json.loads(tool_call.function.arguments)
                
                print(f"🛠️ Calling tool: {tool_name} with args: {tool_args}")
                # Call tool via Klavis API
                tool_result = self.klavis.call_tool(
                    server_url=self.mcp_server_url,
                    tool_name=tool_name,
                    tool_args=tool_args
                )
                
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": json.dumps(tool_result)
                })
            
            # 4. Get final response from LLM
            final_response = self.fireworks.chat.completions.create(
                model=self.model,
                messages=messages
            )
            return final_response.choices[0].message.content
        
        # If no tools needed, return the assistant message directly
        return assistant_message.content