from mcp.server.fastmcp import FastMCP
import httpx
import os
import asyncio
import base64
import json

mcp = FastMCP("Lola Suite MCP Bridge")

BASE_URL = os.getenv("LOLA_API_BASE_URL")
ACCESS_TOKEN = os.getenv("LOLA_ACCESS_TOKEN")
AUTH_TOKEN = os.getenv("LOLA_AUTH_TOKEN")

# -------- Safety checks --------

if not BASE_URL or not BASE_URL.startswith("http"):
    raise RuntimeError("LOLA_API_BASE_URL invalid")

if not ACCESS_TOKEN:
    raise RuntimeError("LOLA_ACCESS_TOKEN required")

if not AUTH_TOKEN:
    raise RuntimeError("LOLA_AUTH_TOKEN required")


def headers():
    return {
        "Authorization": f"Bearer {AUTH_TOKEN}",
        "Content-Type": "application/json",
    }


# -------- Decode allowed tools from token --------

def allowed_tools():
    payload = ACCESS_TOKEN.split(".")[0]
    decoded = base64.b64decode(payload).decode()
    data = json.loads(decoded)
    return set(data["selections"]["lola_server"])


# -------- Dynamic tool registration --------

async def register_tools():
    allowed = allowed_tools()

    async with httpx.AsyncClient() as client:
        res = await client.get(f"{BASE_URL}/tools_lola", headers=headers())
        res.raise_for_status()
        tools = res.json()["tools"]

    for tool in tools:
        name = tool["name"]

        # Only expose tools allowed by token
        if name not in allowed:
            continue

        desc = tool.get("description", "")

        # FastMCP derives schema from function signature
        async def dynamic_tool(arguments: dict, name=name):
            async with httpx.AsyncClient() as client:
                r = await client.post(
                    f"{BASE_URL}/call_lola",
                    json={"name": name, "arguments": arguments},
                    headers=headers(),
                )
                r.raise_for_status()
                return r.json()

        mcp.tool(
            name=name,
            description=desc,
        )(dynamic_tool)


# -------- Startup --------

async def startup():
    await register_tools()


def main():
    asyncio.run(startup())
    mcp.run()


if __name__ == "__main__":
    main()
