#!/usr/bin/env python3
"""
MCP Server for Copilot CLI Integration
Allows Copilot CLI to communicate with Local Agent API
"""

import json
import sys
import requests
from typing import Any, Dict, List, Optional
from urllib.parse import urljoin

LOCAL_AGENT_URL = "http://localhost:8888"
TIMEOUT = 30

class MCPServer:
    """Model Context Protocol Server for Copilot CLI"""
    
    def __init__(self):
        self.methods = {
            "tools/list": self.list_tools,
            "tools/call": self.call_tool,
            "resources/read": self.read_resource,
            "resources/list": self.list_resources,
        }
    
    def list_tools(self, params: Dict = None) -> Dict:
        """List available tools"""
        return {
            "tools": [
                {
                    "name": "local_execute",
                    "description": "Execute shell command on local agent",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "command": {"type": "string", "description": "Shell command to execute"},
                            "timeout": {"type": "integer", "default": 300}
                        },
                        "required": ["command"]
                    }
                },
                {
                    "name": "local_deepseek",
                    "description": "Query DeepSeek model on local agent",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "query": {"type": "string", "description": "Question for DeepSeek"},
                            "model": {"type": "string", "default": "deepseek-chat"},
                            "reasoning": {"type": "boolean", "default": False}
                        },
                        "required": ["query"]
                    }
                },
                {
                    "name": "local_ansible",
                    "description": "Run Ansible playbook on local agent",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "playbook_path": {"type": "string", "description": "Path to playbook"},
                            "extra_vars": {"type": "object", "description": "Extra variables"},
                            "check": {"type": "boolean", "default": False, "description": "Dry run"}
                        },
                        "required": ["playbook_path"]
                    }
                },
                {
                    "name": "local_services",
                    "description": "Get status of local services",
                    "inputSchema": {"type": "object"}
                },
                {
                    "name": "local_logs",
                    "description": "Get recent agent logs",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "lines": {"type": "integer", "default": 50}
                        }
                    }
                }
            ]
        }
    
    def list_resources(self, params: Dict = None) -> Dict:
        """List available resources"""
        return {
            "resources": [
                {
                    "uri": "local-agent://status",
                    "name": "Local Agent Status",
                    "description": "Current status of local agent and services"
                },
                {
                    "uri": "local-agent://logs",
                    "name": "Agent Logs",
                    "description": "Recent agent activity logs"
                },
                {
                    "uri": "local-agent://config",
                    "name": "Agent Configuration",
                    "description": "Local agent configuration"
                }
            ]
        }
    
    def read_resource(self, params: Dict) -> Dict:
        """Read resource"""
        uri = params.get("uri", "")
        
        if uri == "local-agent://status":
            try:
                resp = requests.get(f"{LOCAL_AGENT_URL}/services", timeout=TIMEOUT)
                return {"contents": [{"text": json.dumps(resp.json(), indent=2)}]}
            except Exception as e:
                return {"error": str(e)}
        
        elif uri == "local-agent://logs":
            try:
                resp = requests.get(f"{LOCAL_AGENT_URL}/logs?lines=20", timeout=TIMEOUT)
                logs = "\n".join(resp.json().get("logs", []))
                return {"contents": [{"text": logs}]}
            except Exception as e:
                return {"error": str(e)}
        
        elif uri == "local-agent://config":
            try:
                with open("/opt/local-agent/config/agent.conf") as f:
                    return {"contents": [{"text": f.read()}]}
            except Exception as e:
                return {"error": str(e)}
        
        return {"error": f"Unknown resource: {uri}"}
    
    def call_tool(self, params: Dict) -> Dict:
        """Execute tool"""
        tool_name = params.get("name", "")
        tool_params = params.get("params", {})
        
        try:
            if tool_name == "local_execute":
                resp = requests.post(
                    f"{LOCAL_AGENT_URL}/execute",
                    json=tool_params,
                    timeout=tool_params.get("timeout", 300) + 10
                )
                return {"result": resp.json()}
            
            elif tool_name == "local_deepseek":
                resp = requests.post(
                    f"{LOCAL_AGENT_URL}/deepseek",
                    json=tool_params,
                    timeout=TIMEOUT
                )
                return {"result": resp.json()}
            
            elif tool_name == "local_ansible":
                resp = requests.post(
                    f"{LOCAL_AGENT_URL}/ansible",
                    json=tool_params,
                    timeout=600 + 10
                )
                return {"result": resp.json()}
            
            elif tool_name == "local_services":
                resp = requests.get(f"{LOCAL_AGENT_URL}/services", timeout=TIMEOUT)
                return {"result": resp.json()}
            
            elif tool_name == "local_logs":
                lines = tool_params.get("lines", 50)
                resp = requests.get(f"{LOCAL_AGENT_URL}/logs?lines={lines}", timeout=TIMEOUT)
                return {"result": resp.json()}
            
            else:
                return {"error": f"Unknown tool: {tool_name}"}
        
        except requests.exceptions.ConnectionError:
            return {"error": "Local agent API not available"}
        except Exception as e:
            return {"error": str(e)}
    
    def handle_request(self, request: Dict) -> Dict:
        """Handle MCP request"""
        method = request.get("method", "")
        params = request.get("params", {})
        
        if method in self.methods:
            return self.methods[method](params)
        
        return {"error": f"Unknown method: {method}"}
    
    def run(self):
        """Run MCP server"""
        while True:
            try:
                line = input()
                request = json.loads(line)
                response = self.handle_request(request)
                print(json.dumps(response))
                sys.stdout.flush()
            except EOFError:
                break
            except json.JSONDecodeError as e:
                print(json.dumps({"error": f"Invalid JSON: {str(e)}"}))
                sys.stdout.flush()
            except Exception as e:
                print(json.dumps({"error": str(e)}))
                sys.stdout.flush()

if __name__ == "__main__":
    server = MCPServer()
    server.run()