import asyncio
import json
import os
import sys
import time
import base64
from typing import Optional, List

NAPCT_WS_URL = 'ws://127.0.0.1:3001'
RECEIVE_TIMEOUT = 30

_ws = None
_msg_queue: List[dict] = []


def is_image_file(path: str) -> bool:
    ext = os.path.splitext(path)[1].lower()
    return ext in ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp')


def is_video_file(path: str) -> bool:
    ext = os.path.splitext(path)[1].lower()
    return ext in ('.mp4', '.mov', '.avi', '.mkv', '.flv', '.wmv', '.webm')


def build_message_segments(text: str, file_paths: Optional[List[str]] = None) -> list:
    """构建 NapCat 消息段数组"""
    segments = []
    if text:
        segments.append({"type": "text", "data": {"text": text}})
    if file_paths:
        for path in file_paths:
            abs_path = os.path.abspath(path)
            if not os.path.exists(abs_path):
                raise Exception(f"文件不存在: {abs_path}")
            if is_image_file(abs_path):
                segments.append({"type": "image", "data": {"file": f"file://{abs_path}"}})
            elif is_video_file(abs_path):
                segments.append({"type": "video", "data": {"file": f"file://{abs_path}"}})
            else:
                segments.append({"type": "file", "data": {"file": f"file://{abs_path}"}})
    if not segments:
        segments.append({"type": "text", "data": {"text": ""}})
    return segments


async def ensure_ws():
    global _ws
    if _ws is None or _ws.closed:
        import websockets
        _ws = await websockets.connect(NAPCT_WS_URL)
    return _ws


async def ws_send_action(action: str, params: dict) -> dict:
    """发送 action 到 NapCat 并等待响应"""
    ws = await ensure_ws()
    echo = f"mcp_{int(time.time() * 1000)}"
    payload = {"action": action, "params": params, "echo": echo}
    await ws.send(json.dumps(payload))
    # 等待对应的响应
    timeout_ts = time.time() + 15
    while time.time() < timeout_ts:
        try:
            raw = await asyncio.wait_for(ws.recv(), timeout=2)
            resp = json.loads(raw)
            if resp.get("echo") == echo:
                return resp
            # 不是我们的响应，放入消息队列
            _msg_queue.append(resp)
        except asyncio.TimeoutError:
            continue
    raise Exception("发送消息超时")


async def send_private_message(user_id: int, text: str = "", file_paths: Optional[List[str]] = None) -> dict:
    """发送私聊消息，可附带图片/视频"""
    segments = build_message_segments(text, file_paths)
    return await ws_send_action("send_private_msg", {
        "user_id": user_id,
        "message": segments
    })


async def send_group_message(group_id: int, text: str = "", file_paths: Optional[List[str]] = None) -> dict:
    """发送群聊消息，可附带图片/视频"""
    segments = build_message_segments(text, file_paths)
    return await ws_send_action("send_group_msg", {
        "group_id": group_id,
        "message": segments
    })


async def receive_messages(count: int = 1, user_id: Optional[int] = None,
                           group_id: Optional[int] = None, timeout: int = RECEIVE_TIMEOUT) -> list:
    """接收 N 条消息，可按 user_id 或 group_id 过滤"""
    ws = await ensure_ws()
    results = []
    deadline = time.time() + timeout
    while len(results) < count and time.time() < deadline:
        remain = deadline - time.time()
        if remain <= 0:
            break
        try:
            raw = await asyncio.wait_for(ws.recv(), timeout=min(remain, 2))
            msg = json.loads(raw)
            if msg.get("post_type") != "message":
                continue
            msg_user_id = msg.get("user_id")
            msg_group_id = msg.get("group_id") if msg.get("message_type") == "group" else None
            if user_id is not None and msg_user_id != user_id:
                continue
            if group_id is not None and msg_group_id != group_id:
                continue

            # 提取纯文本和媒体信息
            segments = msg.get("message", [])
            text_parts = []
            media_info = []
            for seg in segments:
                if seg.get("type") == "text":
                    text_parts.append(seg.get("data", {}).get("text", ""))
                elif seg.get("type") in ("image", "video"):
                    media_info.append({
                        "type": seg.get("type"),
                        "url": seg.get("data", {}).get("url", ""),
                        "file": seg.get("data", {}).get("file", "")
                    })

            results.append({
                "message_id": msg.get("message_id"),
                "user_id": msg_user_id,
                "group_id": msg_group_id,
                "message_type": msg.get("message_type"),
                "text": "".join(text_parts).strip(),
                "media": media_info,
                "raw": msg
            })
            _msg_queue.clear()
        except asyncio.TimeoutError:
            continue

    return results


# ============================================================
# 同步包装器（供 stdio MCP 调用）
# ============================================================

_loop = None


def get_loop():
    global _loop
    if _loop is None or _loop.is_closed():
        _loop = asyncio.new_event_loop()
        asyncio.set_event_loop(_loop)
    return _loop


def sync_send_private_message(user_id: int, text: str = "", file_paths: list = None):
    loop = get_loop()
    result = loop.run_until_complete(send_private_message(user_id, text, file_paths or []))
    loop.run_until_complete(_ws.close())
    return result


def sync_send_group_message(group_id: int, text: str = "", file_paths: list = None):
    loop = get_loop()
    result = loop.run_until_complete(send_group_message(group_id, text, file_paths or []))
    loop.run_until_complete(_ws.close())
    return result


def sync_receive_messages(count: int = 1, user_id: int = None,
                          group_id: int = None, timeout: int = RECEIVE_TIMEOUT):
    loop = get_loop()
    result = loop.run_until_complete(receive_messages(count, user_id, group_id, timeout))
    loop.run_until_complete(_ws.close())
    return result


# ============================================================
# MCP JSON-RPC 处理
# ============================================================

TOOLS = [
    {
        "name": "send_private_message",
        "description": "向指定 QQ 用户发送私聊消息，可附带任意类型的文件（图片/视频/文档等）。",
        "inputSchema": {
            "type": "object",
            "properties": {
                "user_id": {"type": "number", "description": "目标 QQ 号"},
                "text": {"type": "string", "description": "要发送的文字内容（可为空，但需要至少有文字或媒体文件之一）"},
                "file_paths": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "文件的本地绝对路径列表，图片用 image 类型，视频用 video 类型，其余后缀一律用 file 类型发送"
                }
            },
            "required": ["user_id"]
        }
    },
    {
        "name": "send_group_message",
        "description": "向指定 QQ 群发送消息，可附带任意类型的文件（图片/视频/文档等）。",
        "inputSchema": {
            "type": "object",
            "properties": {
                "group_id": {"type": "number", "description": "目标群号"},
                "text": {"type": "string", "description": "要发送的文字内容（可为空，但需要至少有文字或媒体文件之一）"},
                "file_paths": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "文件的本地绝对路径列表，图片用 image 类型，视频用 video 类型，其余后缀一律用 file 类型发送"
                }
            },
            "required": ["group_id"]
        }
    },
    {
        "name": "receive_messages",
        "description": "接收来自 NapCat 的消息，可按 QQ 号或群号过滤，接收指定条数后返回（有超时限制）。",
        "inputSchema": {
            "type": "object",
            "properties": {
                "count": {"type": "number", "description": "要接收的消息条数，默认 1"},
                "user_id": {"type": "number", "description": "按 QQ 号过滤（可选，不传则不过滤）"},
                "group_id": {"type": "number", "description": "按群号过滤（可选，不传则不过滤）"},
                "timeout": {"type": "number", "description": "超时秒数，默认 30"}
            },
            "required": []
        }
    }
]

TOOL_MAP = {
    "send_private_message": sync_send_private_message,
    "send_group_message": sync_send_group_message,
    "receive_messages": sync_receive_messages
}


def handle_request(req):
    method = req.get("method", "")
    req_id = req.get("id")

    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": "ros-qq-messages",
                    "version": "1.0.0"
                }
            }
        }

    if method == "ping":
        return {"jsonrpc": "2.0", "id": req_id, "result": {}}

    if method == "notifications/initialized":
        return None

    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}}

    if method == "tools/call":
        tool_name = req["params"]["name"]
        arguments = req["params"].get("arguments", {})
        tool_func = TOOL_MAP.get(tool_name)
        if not tool_func:
            return {"jsonrpc": "2.0", "id": req_id,
                    "error": {"code": -32601, "message": f"未知工具: {tool_name}"}}
        try:
            result = tool_func(**arguments)
            text = json.dumps(result, ensure_ascii=False, indent=2)
            return {"jsonrpc": "2.0", "id": req_id,
                    "result": {"content": [{"type": "text", "text": text}]}}
        except Exception as e:
            return {"jsonrpc": "2.0", "id": req_id,
                    "error": {"code": -32000, "message": str(e)}}

    return {"jsonrpc": "2.0", "id": req_id,
            "error": {"code": -32601, "message": f"未知方法: {method}"}}


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            resp = handle_request(req)
            if resp is None:
                continue
            sys.stdout.write(json.dumps(resp, ensure_ascii=False) + '\n')
            sys.stdout.flush()
        except json.JSONDecodeError:
            pass


if __name__ == '__main__':
    main()