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_BASE_URL")
ACCESS_TOKEN = os.getenv("LOLA_ACCESS_TOKEN")


# ---------------- SAFETY CHECK ----------------

if not BASE_URL or not BASE_URL.startswith("http"):
    raise RuntimeError("LOLA_BASE_URL must be a full https:// URL")

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


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


# ---------------- DECODE TOKEN ----------------

def get_allowed_tools():
    """
    Extract allowed tool list from the access token payload
    """
    payload = ACCESS_TOKEN.split(".")[0]
    decoded = base64.b64decode(payload).decode()
    data = json.loads(decoded)

    return set(
        data.get("selections", {})
            .get("lola_server", [])
    )


# ---------------- DYNAMIC TOOL REGISTRATION ----------------

async def register_tools():
    allowed = get_allowed_tools()

    async with httpx.AsyncClient(timeout=30.0) 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:
        tool_name = tool["name"]

        # 🔒 Filter tools by access token
        if tool_name not in allowed:
            continue

        input_schema = tool["input_schema"]
        description = tool.get("description", "")

        # ⚠️ Important: capture tool_name correctly (closure fix)
        async def dynamic_tool(tool_name=tool_name, **kwargs):
            async with httpx.AsyncClient(timeout=30.0) as client:
                r = await client.post(
                    f"{BASE_URL}/call_lola",
                    json={"name": tool_name, "arguments": kwargs},
                    headers=headers(),
                )
                r.raise_for_status()
                return r.json()

        mcp.tool(
            name=tool_name,
            description=description,
            input_schema=input_schema,
        )(dynamic_tool)


# ---------------- MAIN ----------------

async def startup():
    await register_tools()


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


if __name__ == "__main__":
    main()
