#!/usr/bin/env python3
"""
Base MCP server implementation - shared between standard and dynamic servers.
"""

import json
import os
import sys
import time

import requests


class BaseMCPServer:
    """Base MCP server that handles JSON-RPC directly."""

    def __init__(self):
        self.tools = []
        self.bundle_name = os.getenv("BUNDLE_NAME")
        self.load_tools()

    def load_tools(self):
        """Load tools from gateway."""
        try:
            gateway_url = os.getenv("GATEWAY_URL", "http://localhost:8000")

            # Use bundle-specific endpoint if bundle is specified
            if self.bundle_name:
                endpoint = f"{gateway_url}/mcp/tools/{self.bundle_name}/quick"
            else:
                endpoint = f"{gateway_url}/mcp/tools/quick"

            print(f"Loading tools from: {endpoint} (attempt 1/3)", file=sys.stderr)
            
            # Try to connect with retries
            for attempt in range(3):
                try:
                    response = requests.get(endpoint, timeout=5)
                    print(f"Response received: status={response.status_code}", file=sys.stderr)
                    
                    if response.status_code == 200:
                        data = response.json()
                        self.tools = data.get("tools", [])
                        bundle_info = f" (bundle: {self.bundle_name})" if self.bundle_name else " (all bundles)"
                        print(f"Loaded {len(self.tools)} tools{bundle_info}", file=sys.stderr)
                        break
                    else:
                        print(f"Failed to load tools: {response.status_code}", file=sys.stderr)
                        self.tools = []
                        
                    break
                except requests.exceptions.RequestException as e:
                    print(f"Request error (attempt {attempt + 1}/3): {e}", file=sys.stderr)
                    if attempt < 2:
                        time.sleep(2)
                        print(f"Loading tools from: {endpoint} (attempt {attempt + 2}/3)", file=sys.stderr)
                    else:
                        print("Failed to connect to gateway after 3 attempts", file=sys.stderr)
                        self.tools = []
                        
        except Exception as e:
            print(f"Error loading tools: {e}", file=sys.stderr)
            self.tools = []

    def handle_initialize(self, request_id):
        """Handle initialize request."""
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "protocolVersion": "0.1.0",
                "capabilities": {
                    "tools": {},
                    "resources": {"subscribe": False},
                    "prompts": {}
                },
                "serverInfo": {
                    "name": "coherence-mcp-server",
                    "version": "0.1.0"
                }
            }
        }

    def handle_tools_list(self, request_id):
        """Handle tools/list request."""
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "tools": self.tools
            }
        }

    def handle_tools_call(self, request_id, params):
        """Handle tools/call request."""
        tool_name = params.get("name")
        arguments = params.get("arguments", {})

        # Find the tool
        tool = None
        for t in self.tools:
            if t.get("name") == tool_name:
                tool = t
                break

        if not tool:
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32602,
                    "message": f"Tool not found: {tool_name}"
                }
            }

        # Call the gateway to execute the tool
        try:
            gateway_url = os.getenv("GATEWAY_URL", "http://localhost:8000")
            jwt_token = os.getenv("JWT_TOKEN", "")

            headers = {}
            if jwt_token:
                headers["Authorization"] = f"Bearer {jwt_token}"

            # Use bundle-specific endpoint if bundle is specified
            if self.bundle_name:
                endpoint = f"{gateway_url}/mcp/tools/{self.bundle_name}/call"
            else:
                endpoint = f"{gateway_url}/mcp/tools/call"

            response = requests.post(
                endpoint,
                json={
                    "name": tool_name,
                    "arguments": arguments
                },
                headers=headers,
                timeout=30
            )

            if response.status_code == 200:
                result = response.json()
                return {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": result
                }
            else:
                return {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "error": {
                        "code": -32603,
                        "message": f"Tool execution failed: {response.status_code}"
                    }
                }

        except Exception as e:
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32603,
                    "message": str(e)
                }
            }

    def handle_resources_list(self, request_id):
        """Handle resources/list request."""
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "resources": []
            }
        }

    def handle_resources_read(self, request_id, params):
        """Handle resources/read request."""
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32601,
                "message": "Resources not implemented"
            }
        }

    def handle_prompts_list(self, request_id):
        """Handle prompts/list request."""
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "prompts": []
            }
        }

    def run(self):
        """Run the server."""
        # In Docker mode, keep the server running for health checks
        if os.environ.get('MCP_MODE') == 'docker':
            print("MCP Server running in Docker mode - listening on stdin...", file=sys.stderr)
            sys.stderr.flush()
        
        while True:
            try:
                # Read request from stdin
                line = sys.stdin.readline()
                if not line:
                    # In Docker mode, sleep and continue instead of breaking
                    if os.environ.get('MCP_MODE') == 'docker':
                        time.sleep(1)
                        continue
                    break

                request = json.loads(line.strip())
                method = request.get("method")
                request_id = request.get("id")

                if method == "initialize":
                    response = self.handle_initialize(request_id)
                elif method == "tools/list":
                    response = self.handle_tools_list(request_id)
                elif method == "tools/call":
                    response = self.handle_tools_call(request_id, request.get("params", {}))
                elif method == "resources/list":
                    response = self.handle_resources_list(request_id)
                elif method == "resources/read":
                    response = self.handle_resources_read(request_id, request.get("params", {}))
                elif method == "prompts/list":
                    response = self.handle_prompts_list(request_id)
                else:
                    response = {
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "error": {"code": -32601, "message": "Method not found"}
                    }

                # Send response to stdout
                sys.stdout.write(json.dumps(response) + "\n")
                sys.stdout.flush()

            except Exception as e:
                # Send error response
                error_response = {
                    "jsonrpc": "2.0",
                    "id": request.get("id") if 'request' in locals() else None,
                    "error": {"code": -32603, "message": str(e)}
                }
                sys.stdout.write(json.dumps(error_response) + "\n")
                sys.stdout.flush()