from mcp.server.fastmcp import FastMCP
from pymongo import MongoClient
from datetime import datetime
import os

# Any agent (Claude, GPT-4o, Llama) can now query your MongoDB
mcp = FastMCP("zetwerk-inventory-server")
db  = MongoClient(os.getenv("MONGO_URI"))["production"]

# ── TOOLS (agent can call these) ──────────────────────────────

@mcp.tool()
async def search_inventory(
    product_name: str,
    min_stock: int = 0,
    category: str = None
) -> list[dict]:
    """
    Search the product inventory by name and optional filters.
    Returns products with stock levels, pricing, and supplier info.
    Use when an agent needs to check product availability.
    Args:
        product_name: Full or partial product name to search
        min_stock: Minimum stock quantity (default 0 = include out-of-stock)
        category: Optional category filter (e.g. 'fasteners', 'steel')
    """
    query = {"name": {"$regex": product_name, "$options": "i"}}
    if min_stock: query["stock"] = {"$gte": min_stock}
    if category:  query["category"] = category
    return list(db.inventory.find(query, {"_id": 0}).limit(20))

@mcp.tool()
async def get_supplier_quotes(product_id: str) -> list[dict]:
    """Get all supplier quotes for a specific product ID, sorted by price."""
    return list(db.quotes.find(
        {"product_id": product_id},
        {"_id": 0}
    ).sort("unit_price", 1))

@mcp.tool()
async def create_purchase_order(
    supplier_id: str,
    product_id: str,
    quantity: int,
    requested_by: str
) -> dict:
    """
    Create a purchase order. REQUIRES human approval — do not call autonomously.
    This creates a real order in the production system.
    """
    order = {
        "supplier_id": supplier_id, "product_id": product_id,
        "quantity": quantity, "requested_by": requested_by,
        "status": "pending_approval", "created_at": datetime.utcnow()
    }
    result = db.purchase_orders.insert_one(order)
    return {"order_id": str(result.inserted_id), "status": "pending_approval"}

# ── RESOURCES (agent can read these for context) ──────────────

@mcp.resource("inventory://schema")
async def get_schema() -> str:
    """Database schema — agent reads this to understand data structure"""
    return """
    Collections:
    - inventory: {product_id, name, category, stock, unit, min_reorder_qty}
    - quotes: {product_id, supplier_id, unit_price, lead_time_days, valid_until}
    - suppliers: {supplier_id, name, rating, location, payment_terms}
    - purchase_orders: {order_id, supplier_id, product_id, quantity, status}
    """

if __name__ == "__main__":
    mcp.run(transport="stdio")  # Local process mode
    # Or: mcp.run(transport="sse", port=8080)  # Remote HTTP mode

import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { MongoClient } from "mongodb";

const server = new Server(
  { name: "postgres-agent-server", version: "1.0.0" },
  { capabilities: { tools: {}, resources: {} } }
);

server.setRequestHandler("tools/list", async () => ({
  tools: [{
    name: "run_query",
    description: "Execute a read-only SQL query on the analytics database. SELECT only — no mutations.",
    inputSchema: {
      type: "object",
      properties: {
        sql: { type: "string", description: "SELECT query to run" },
        limit: { type: "number", default: 100 }
      },
      required: ["sql"]
    }
  }]
}));

server.setRequestHandler("tools/call", async (request) => {
  const { name, arguments: args } = request.params;
  if (name === "run_query") {
    // Safety: reject any non-SELECT query
    if (!args.sql.trim().toUpperCase().startsWith("SELECT"))
      throw new Error("Only SELECT queries are allowed");
    const rows = await pgClient.query(args.sql + ` LIMIT ${args.limit || 100}`);
    return { content: [{ type: "text", text: JSON.stringify(rows.rows, null, 2) }] };
  }
});

const transport = new StdioServerTransport();
await server.connect(transport);

from mcp.server.fastmcp import FastMCP
from mcp.server.auth import OAuthProvider
import jwt

mcp = FastMCP("secure-data-server")

def verify_token(token: str) -> dict:
    """Verify JWT from OAuth 2.1 flow"""
    try:
        payload = jwt.decode(token, PUBLIC_KEY, algorithms=["RS256"])
        return {"user_id": payload["sub"], "scopes": payload.get("scope", []).split()}
    except jwt.InvalidTokenError:
        raise PermissionError("Invalid or expired token")

@mcp.tool()
async def get_user_data(user_id: str, ctx: Context) -> dict:
    """Get user account data — requires 'read:users' scope"""
    claims = verify_token(ctx.auth_token)
    if "read:users" not in claims["scopes"]:
        raise PermissionError("Missing required scope: read:users")
    return db.users.find_one({"id": user_id})

import httpx, uuid, asyncio

class A2AClient:
    """Send tasks to remote A2A-compatible agents"""

    def __init__(self, agent_url: str, auth_token: str):
        self.agent_url = agent_url.rstrip("/")
        self.headers = {"Authorization": f"Bearer {auth_token}"}

    async def discover(self) -> dict:
        """Read the agent's capabilities from its Agent Card"""
        async with httpx.AsyncClient() as client:
            r = await client.get(f"{self.agent_url}/.well-known/agent.json")
            return r.json()

    async def send_task(
        self,
        message: str,
        webhook_url: str = None
    ) -> dict:
        """Submit a task to the remote agent"""
        task_id = str(uuid.uuid4())
        payload = {
            "id": task_id,
            "message": {"role": "user", "parts": [{"text": message}]},
            "webhookUrl": webhook_url
        }
        async with httpx.AsyncClient() as client:
            r = await client.post(
                f"{self.agent_url}/tasks/send",
                json=payload, headers=self.headers
            )
            return r.json()

    async def wait_for_completion(self, task_id: str, timeout: int = 300) -> dict:
        """Poll until task completes (use webhooks in production instead)"""
        async with httpx.AsyncClient() as client:
            for _ in range(timeout // 5):
                r = await client.get(
                    f"{self.agent_url}/tasks/{task_id}",
                    headers=self.headers
                )
                task = r.json()
                if task["status"]["state"] in ("completed", "failed"):
                    return task
                await asyncio.sleep(5)
        raise TimeoutError(f"Task {task_id} did not complete in {timeout}s")

# Usage — orchestrator delegates to specialist agents
async def run_investment_research(company: str):
    research_client  = A2AClient("https://research.agentco.com",  token)
    financial_client = A2AClient("https://finance.agentco.com",   token)

    # Discover capabilities first
    research_card  = await research_client.discover()
    financial_card = await financial_client.discover()

    # Run both in parallel
    research_task, financial_task = await asyncio.gather(
        research_client.send_task(f"Research {company} — find all public information"),
        financial_client.send_task(f"Analyse {company} financials and estimate valuation")
    )
    # Wait for both
    research_result, financial_result = await asyncio.gather(
        research_client.wait_for_completion(research_task["id"]),
        financial_client.wait_for_completion(financial_task["id"])
    )
    return {"research": research_result, "financials": financial_result}