"""MCP Streamable-HTTP client for opax-mcp."""
import json
import os
import httpx

OPAX_MCP_URL = os.environ.get(
    "MCP_SERVER_URL",
    "https://opax-mcp-357036551735.us-central1.run.app",
).rstrip("/")
MCP_SECRET = os.environ.get("MCP_SECRET", "")

def _headers(session_id=None):
    headers = {
        "Authorization": f"Bearer {MCP_SECRET}",
        "Content-Type": "application/json",
        "Accept": "application/json, text/event-stream",
        "MCP-Protocol-Version": "2025-06-18",
    }
    if session_id:
        headers["Mcp-Session-Id"] = session_id
    return headers

def _decode(response):
    raw = response.text
    data = [line[5:].strip() for line in raw.splitlines() if line.startswith("data:")]
    return json.loads(data[-1] if data else raw)

def _post(payload, session_id=None):
    with httpx.Client(timeout=60) as client:
        response = client.post(OPAX_MCP_URL, headers=_headers(session_id), json=payload)
        response.raise_for_status()
        return _decode(response), response.headers.get("Mcp-Session-Id")

def _session():
    _, session_id = _post({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "initialize",
        "params": {
            "protocolVersion": "2025-06-18",
            "capabilities": {},
            "clientInfo": {"name": "osvxcc", "version": "1.0"},
        },
    })
    _post({"jsonrpc": "2.0", "method": "notifications/initialized"}, session_id)
    return session_id

def list_tools():
    session_id = _session()
    response, _ = _post(
        {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}},
        session_id,
    )
    return response.get("result", {}).get("tools", [])

def call_tool(name, arguments=None):
    session_id = _session()
    response, _ = _post(
        {
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {"name": name, "arguments": arguments or {}},
        },
        session_id,
    )
    return response.get("result", response)

def run_jason(prompt: str) -> dict:
    return call_tool("run_jason", {"prompt": prompt})

def run_emma(prompt: str) -> dict:
    return call_tool("run_emma", {"prompt": prompt})