#!/usr/bin/env python3
"""Wolfram Alpha MCP Server — AI agent interface to Wolfram Alpha."""
import json, sys, os, urllib.request, urllib.parse, urllib.error

APP_ID = os.environ.get("WOLFRAM_APP_ID", "")
BASE = "https://api.wolframalpha.com/v2"

def query_wolfram(query: str) -> str:
    """Query Wolfram Alpha and return plaintext result."""
    if not APP_ID:
        return "WOLFRAM_APP_ID not set"
    params = urllib.parse.urlencode({"input": query, "appid": APP_ID, "format": "plaintext"})
    try:
        with urllib.request.urlopen(f"{BASE}/query?{params}", timeout=15) as r:
            return r.read().decode()
    except Exception as e:
        return f"Error: {e}"

def handle_request(req):
    method = req.get("method", "")
    params = req.get("params", {})

    if method == "list_tools":
        return {"tools": [
            {"name": "wolfram_query", "description": "Query Wolfram Alpha for computational answers",
             "inputSchema": {"type": "object", "properties": {
                 "query": {"type": "string", "description": "Query in natural language or Wolfram syntax"}
             }}},
            {"name": "wolfram_math", "description": "Solve a mathematical expression",
             "inputSchema": {"type": "object", "properties": {
                 "expression": {"type": "string", "description": "Mathematical expression"}
             }}},
        ]}
    elif method == "call_tool":
        name = params.get("name", "")
        args = params.get("arguments", {})
        if name == "wolfram_query":
            result = query_wolfram(args.get("query", ""))
        elif name == "wolfram_math":
            result = query_wolfram(args.get("expression", ""))
        else:
            result = f"unknown tool: {name}"
        return {"content": [{"type": "text", "text": result}]}
    elif method == "list_resource_templates":
        return {"resourceTemplates": []}
    return {"error": f"unknown method: {method}"}

def main():
    while True:
        line = sys.stdin.readline()
        if not line: break
        try:
            req = json.loads(line)
            resp = handle_request(req)
            sys.stdout.write(json.dumps(resp) + "\n")
            sys.stdout.flush()
        except Exception as e:
            sys.stdout.write(json.dumps({"error": str(e)}) + "\n")
            sys.stdout.flush()

if __name__ == "__main__":
    main()
