from mcp.server.fastmcp import FastMCP
import os
import json
import uuid
import asyncio
import logging
import websockets

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("openclaw_mcp")

mcp = FastMCP("OpenClaw Assistant")

OPENCLAW_HTTP = os.getenv("OPENCLAW_URL", "http://192.168.50.139:18789").rstrip("/")
OPENCLAW_TOKEN = os.getenv("OPENCLAW_TOKEN", "").strip()

OPENCLAW_WS_URLS = [
    OPENCLAW_HTTP.replace("http://", "ws://").replace("https://", "wss://") + "/ws",
    OPENCLAW_HTTP.replace("http://", "ws://").replace("https://", "wss://") + "/gateway",
    OPENCLAW_HTTP.replace("http://", "ws://").replace("https://", "wss://") + "/rpc",
]


async def rpc_call(method: str, params: dict, timeout: int = 90) -> dict:
    last_error = None

    for ws_url in OPENCLAW_WS_URLS:
        try:
            logger.info("Trying OpenClaw WS: %s", ws_url)

            headers = {}
            if OPENCLAW_TOKEN:
                headers["Authorization"] = f"Bearer {OPENCLAW_TOKEN}"

            async with websockets.connect(
                ws_url,
                additional_headers=headers,
                ping_interval=20,
                ping_timeout=20,
            ) as ws:
                request_id = str(uuid.uuid4())

                payload = {
                    "id": request_id,
                    "method": method,
                    "params": params,
                }

                logger.info("Sending RPC: %s", payload)
                await ws.send(json.dumps(payload, ensure_ascii=False))

                while True:
                    raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
                    logger.info("OpenClaw WS recv: %s", raw[:1000])

                    try:
                        data = json.loads(raw)
                    except Exception:
                        continue

                    if data.get("id") == request_id:
                        if "error" in data:
                            return {
                                "success": False,
                                "message": str(data["error"])[:900],
                            }

                        return {
                            "success": True,
                            "message": json.dumps(
                                data.get("result", data),
                                ensure_ascii=False,
                            )[:900],
                        }

                    # 有些 OpenClaw 结果可能通过 event 推送
                    if data.get("method") in ["chat.delta", "chat.final", "chat.message"]:
                        return {
                            "success": True,
                            "message": json.dumps(data, ensure_ascii=False)[:900],
                        }

        except Exception as e:
            last_error = e
            logger.exception("WS failed: %s", ws_url)

    return {
        "success": False,
        "message": f"OpenClaw WebSocket 调用失败: {last_error}",
    }


def send_to_openclaw(message: str) -> dict:
    params = {
        "sessionKey": "main",
        "message": message,
        "deliver": False,
        "idempotencyKey": str(uuid.uuid4()),
    }

    return asyncio.run(rpc_call("chat.send", params))


@mcp.tool()
def ask_openclaw(question: str) -> dict:
    """
    Use this tool when the user wants OpenClaw to answer, think, reason,
    control Home Assistant, inspect smart home state, or execute an agent task.
    """
    return send_to_openclaw(question)


@mcp.tool()
def control_home_assistant(command: str) -> dict:
    """
    Use this tool when the user wants to control Home Assistant devices
    through OpenClaw, including lights, switches, curtains, air conditioners,
    scenes, automations, locks, cameras, media players, and sensors.
    """
    return send_to_openclaw(command)


if __name__ == "__main__":
    mcp.run(transport="stdio")