server.py

import os
import json
import httpx
import yaml
from typing import Any, Dict, List, Optional
from mcp.server.fastmcp import FastMCP
from dotenv import set_key, load_dotenv

# Initialize FastMCP server
mcp = FastMCP("API Wizard")

# Paths
ENV_PATH = ".env"
CONNECTIONS_PATH = "connections.json"
HISTORY_PATH = "history.json"
SPECS_DIR = "specs"

# Ensure files and directories exist
os.makedirs(SPECS_DIR, exist_ok=True)
if not os.path.exists(ENV_PATH): open(ENV_PATH, "w").close()
if not os.path.exists(CONNECTIONS_PATH): 
    with open(CONNECTIONS_PATH, "w") as f: json.dump({}, f)
if not os.path.exists(HISTORY_PATH):
    with open(HISTORY_PATH, "w") as f: json.dump([], f)

def get_connections() -> Dict[str, Any]:
    try:
        with open(CONNECTIONS_PATH, "r") as f: return json.load(f)
    except: return {}

def save_connections(connections: Dict[str, Any]):
    with open(CONNECTIONS_PATH, "w") as f: json.dump(connections, f, indent=2)

def log_request(name: str, method: str, path: str, status: int):
    try:
        with open(HISTORY_PATH, "r") as f: history = json.load(f)
    except: history = []
    history.insert(0, {"connection": name, "method": method, "path": path, "status": status})
    with open(HISTORY_PATH, "w") as f: json.dump(history[:50], f, indent=2)

@mcp.tool()
def set_secret(key_name: str, value: str) -> str:
    """Securely saves a credential."""
    set_key(ENV_PATH, key_name, value)
    return f"Secret '{key_name}' saved."

@mcp.tool()
def save_connection(name: str, base_url: str, secret_key_name: str, auth_type: str = "Bearer") -> str:
    """Saves an API connection profile."""
    connections = get_connections()
    connections[name] = {"base_url": base_url, "secret_key_name": secret_key_name, "auth_type": auth_type}
    save_connections(connections)
    return f"Connection '{name}' saved."

@mcp.tool()
async def execute_request(connection_name: str, endpoint_path: str, method: str = "GET", json_data: dict = None, params: dict = None) -> Dict[str, Any]:
    """Executes an API request using a saved profile."""
    connections = get_connections()
    if connection_name not in connections: return {"error": "Not found"}
    profile = connections[connection_name]
    load_dotenv(ENV_PATH, override=True)
    secret = os.getenv(profile["secret_key_name"])
    if not secret: return {"error": "Secret missing"}
    
    headers = {"Authorization": f"Bearer {secret}"} if profile["auth_type"] == "Bearer" else {"X-API-Key": secret}
    url = f"{profile['base_url'].rstrip('/')}/{endpoint_path.lstrip('/')}"
    
    async with httpx.AsyncClient() as client:
        try:
            response = await client.request(method=method, url=url, headers=headers, json=json_data, params=params, timeout=30.0)
            log_request(connection_name, method, endpoint_path, response.status_code)
            return {"status": response.status_code, "data": response.json() if "json" in response.headers.get("content-type", "") else response.text}
        except Exception as e:
            return {"error": str(e)}

@mcp.tool()
async def load_openapi_spec(url: str, name: str) -> str:
    """Downloads an OpenAPI spec (JSON/YAML) and registers it."""
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(url)
            content = response.text
            # Try parsing as JSON first, then YAML
            try: spec = json.loads(content)
            except: spec = yaml.safe_load(content)
            
            path = os.path.join(SPECS_DIR, f"{name}.json")
            with open(path, "w") as f: json.dump(spec, f, indent=2)
            return f"Spec '{name}' loaded and saved to {path}"
        except Exception as e:
            return f"Failed to load spec: {str(e)}"

@mcp.tool()
def list_spec_endpoints(api_name: str) -> List[str]:
    """Lists all available endpoints and methods from a loaded spec."""
    path = os.path.join(SPECS_DIR, f"{api_name}.json")
    if not os.path.exists(path): return [f"Spec '{api_name}' not found."]
    
    with open(path, "r") as f: spec = json.load(f)
    endpoints = []
    for p, methods in spec.get("paths", {}).items():
        for m in methods.keys():
            endpoints.append(f"{m.upper()} {p}")
    return sorted(endpoints)

@mcp.tool()
async def execute_spec_call(api_name: str, connection_name: str, path: str, method: str, params: dict = None, body: dict = None) -> Dict[str, Any]:
    """Executes a call by looking up the endpoint in the spec."""
    # This currently bridges the spec knowledge with existing execute_request
    return await execute_request(connection_name, path, method, json_data=body, params=params)

@mcp.tool()
def export_postman_collection() -> str:
    """Generates a Postman collection JSON from saved connections."""
    connections = get_connections()
    items = [{"name": n, "request": {"method": "GET", "url": {"raw": p['base_url']}}} for n, p in connections.items()]
    return json.dumps({"info": {"name": "API Wizard Export", "schema": "https://schema.getpostman.com/json/collection/v2.1.0/collection.json"}, "item": items}, indent=2)

if __name__ == "__main__":
    mcp.run()