"""
DG-Lab Coyote WebSocket 控制插件
基于 DG-LAB-OPENSOURCE 官方 Socket V2 协议
支持郊狼 3.0 脉冲主机的 WebSocket 远程控制

参考文档:
- https://github.com/DG-LAB-OPENSOURCE/DG-LAB-OPENSOURCE
- Socket V2 协议: socket/v2/README.md
"""

import json
import time
import threading
import uuid
import os

try:
    import websocket
    HAS_WS = True
except ImportError:
    HAS_WS = False

try:
    import qrcode
    HAS_QRCODE = True
except ImportError:
    HAS_QRCODE = False


class Liugin:
    """DG-Lab 郊狼 WebSocket 控制插件"""

    def __init__(self):
        self.ws = None
        self.client_id = ""
        self.target_id = ""
        self.connected = False
        self.paired = False
        self.ws_url = ""
        self.strength_a = 0
        self.strength_b = 0
        self.limit_a = 100
        self.limit_b = 100
        self._heartbeat_thread = None
        self._running = False

        self.usage = """DG-Lab 郊狼控制插件 使用说明：

/dglab connect <ws_url>  - 连接到 WebSocket 服务器
/dglab qr                 - 生成配对二维码 (需APP扫码)
/dglab status             - 查看连接状态和通道强度
/dglab strength <ch> <+|-|=|set> <val>  - 控制通道强度
    ch: a/b (通道A/B)
    操作: +增加, -减少, =设置为指定值
    示例: /dglab strength a + 10
          /dglab strength b = 50
/dglab pulse <ch> <freq> <intensity> <duration>  - 发送波形
    ch: a/b
    freq: 频率 (10-240)
    intensity: 强度 (0-100)
    duration: 持续秒数
    示例: /dglab pulse a 30 50 5
/dglab wave <ch> <hex_data>  - 发送原始波形HEX数据
    示例: /dglab wave a 0A0A0A0A64646464
/dglab clear <ch>          - 清空通道波形队列 (a/b/both)
/dglab limit <ch> <val>    - 设置通道强度软上限 (0-200)
/dglab disconnect          - 断开连接
"""

    def set_cli(self, cli):
        """注册 CLI 命令"""
        self.cli = cli
        if hasattr(cli, 'register_liugin_command'):
            cli.register_liugin_command('dglab', self._cmd_handler)

    def _cmd_handler(self, args):
        """处理 /dglab 命令"""
        parts = args.strip().split(maxsplit=1)
        if not parts:
            return self.handle("help")
        return self.handle(args)

    def get_tool_info(self):
        return {
            "name": "dglab_ws",
            "description": "DG-Lab 郊狼 3.0 WebSocket 控制插件，支持连接、配对、强度控制、波形发送",
            "keywords": [
                "dg-lab", "dglab", "郊狼", "coyote", "websocket",
                "脉冲", "强度", "波形", "蓝牙", "控制", "情趣"
            ],
            "usage": self.usage
        }

    def get_mcp_definition(self):
        return {
            "name": "dglab_ws",
            "description": "DG-Lab Coyote 3.0 WebSocket controller plugin",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": [
                            "connect", "disconnect", "status",
                            "strength", "pulse", "wave",
                            "clear", "limit", "qr"
                        ]
                    },
                    "args": {"type": "string", "description": "命令参数"}
                },
                "required": ["action"]
            }
        }

    def convert_mcp_args(self, arguments):
        action = arguments.get("action", "")
        args = arguments.get("args", "")
        return f"{action} {args}"

    # ── 核心协议方法 ──────────────────────────────────────

    def _on_message(self, ws, raw):
        """处理服务端消息"""
        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            return

        msg_type = data.get("type", "")
        msg = data.get("message", "")
        cid = data.get("clientId", "")
        tid = data.get("targetId", "")

        if msg_type == "bind":
            if cid and not self.client_id:
                self.client_id = cid
            if msg == "200":
                self.paired = True
                self.target_id = tid
                print(f"\n[DG-Lab] ✅ 配对成功! targetId={tid}")
            elif msg == "400":
                print("\n[DG-Lab] ❌ 配对失败: 该ID已被绑定")
            else:
                print(f"\n[DG-Lab] 绑定状态: {msg}")

        elif msg_type == "msg":
            if msg.startswith("strength-"):
                self._parse_strength(msg)
            elif msg.startswith("feedback-"):
                btn = msg.split("-")[1]
                print(f"\n[DG-Lab] APP反馈按钮: {btn}")

        elif msg_type == "heartbeat":
            pass

        elif msg_type == "break":
            self.paired = False
            print("\n[DG-Lab] ⚠️ APP端已断开连接")

        elif msg_type == "error":
            print(f"\n[DG-Lab] ❌ 错误: {msg}")

    def _on_error(self, ws, error):
        print(f"\n[DG-Lab] WebSocket错误: {error}")

    def _on_close(self, ws, code, reason):
        self.connected = False
        self.paired = False
        self._running = False
        print(f"\n[DG-Lab] 连接已关闭 (code={code}, reason={reason})")

    def _on_open(self, ws):
        self.connected = True
        print("\n[DG-Lab] ✅ 已连接到WebSocket服务器")
        self._running = True
        self._heartbeat_thread = threading.Thread(
            target=self._heartbeat_loop, daemon=True
        )
        self._heartbeat_thread.start()

    def _heartbeat_loop(self):
        """每60秒发送心跳"""
        while self._running and self.connected:
            time.sleep(60)
            if self.ws and self.connected:
                try:
                    hb = {
                        "type": "heartbeat",
                        "clientId": self.client_id,
                        "targetId": self.target_id,
                        "message": "ping"
                    }
                    self.ws.send(json.dumps(hb))
                except Exception:
                    break

    def _parse_strength(self, msg):
        """解析APP回传的强度: strength-A+B+A上限+B上限"""
        try:
            vals = msg.replace("strength-", "").split("+")
            if len(vals) >= 4:
                self.strength_a = int(vals[0])
                self.strength_b = int(vals[1])
                self.limit_a = int(vals[2])
                self.limit_b = int(vals[3])
                print(
                    f"\n[DG-Lab] 📊 强度状态: "
                    f"A={self.strength_a}/{self.limit_a}  "
                    f"B={self.strength_b}/{self.limit_b}"
                )
        except (ValueError, IndexError):
            pass

    def _send_to_app(self, msg_type, message):
        """发送消息到服务端 (转发给APP)"""
        if not self.connected or not self.ws:
            return False, "未连接到服务器"
        if not self.paired:
            return False, "尚未与APP配对，请先用APP扫码"

        payload = {
            "type": msg_type,
            "clientId": self.client_id,
            "targetId": self.target_id,
            "message": message
        }
        try:
            self.ws.send(json.dumps(payload))
            return True, "ok"
        except Exception as e:
            return False, str(e)

    def _send_strength(self, channel, mode, value):
        """
        发送强度控制指令
        channel: 1=A, 2=B
        mode: 0=减少, 1=增加, 2=设为指定值
        value: 0-200
        """
        msg = f"strength-{channel}+{mode}+{value}"
        return self._send_to_app("msg", msg)

    def _send_pulse(self, channel, freq, intensity, duration=5):
        """生成并发送波形数据"""
        ch = "A" if channel.lower() == "a" else "B"

        def map_freq(f):
            f = max(10, min(1000, f))
            if f <= 100:
                return f
            elif f <= 600:
                return (f - 100) // 5 + 100
            else:
                return (f - 600) // 10 + 200

        freq_byte = map_freq(freq)
        intensity_byte = max(0, min(100, intensity))

        wave_hex = ""
        for _ in range(25 * max(1, duration)):
            wave_hex += f"{freq_byte:02X}{intensity_byte:02X}" * 4

        pulse_data = []
        for i in range(0, len(wave_hex), 16):
            chunk = wave_hex[i:i + 16]
            if len(chunk) == 16:
                pulse_data.append(chunk)

        if not pulse_data:
            pulse_data = [f"{freq_byte:02X}" * 4 + f"{intensity_byte:02X}" * 4]

        batch_size = min(len(pulse_data), 100)
        message = f"{ch}:{json.dumps(pulse_data[:batch_size])}"

        payload = {
            "type": "clientMsg",
            "channel": ch,
            "time": duration,
            "message": message,
            "clientId": self.client_id,
            "targetId": self.target_id
        }

        if not self.connected or not self.ws:
            return False, "未连接到服务器"
        if not self.paired:
            return False, "尚未与APP配对"

        try:
            self.ws.send(json.dumps(payload))
            return True, f"已发送 {batch_size} 条波形数据到通道{ch}，持续{duration}秒"
        except Exception as e:
            return False, str(e)

    def _send_raw_wave(self, channel, hex_data):
        """发送原始HEX波形数据"""
        ch = "A" if channel.lower() == "a" else "B"
        pulse_data = []
        hex_clean = hex_data.replace(" ", "").replace(",", "")
        for i in range(0, len(hex_clean), 16):
            chunk = hex_clean[i:i + 16]
            if len(chunk) == 16:
                pulse_data.append(chunk)

        if not pulse_data:
            return False, "波形数据不足8字节 (需要16个HEX字符)"

        message = f"{ch}:{json.dumps(pulse_data[:100])}"
        payload = {
            "type": "clientMsg",
            "channel": ch,
            "time": len(pulse_data) // 25 + 1,
            "message": message,
            "clientId": self.client_id,
            "targetId": self.target_id
        }

        if not self.connected or not self.ws:
            return False, "未连接到服务器"
        if not self.paired:
            return False, "尚未与APP配对"

        try:
            self.ws.send(json.dumps(payload))
            return True, f"已发送 {len(pulse_data)} 条波形数据到通道{ch}"
        except Exception as e:
            return False, str(e)

    # ── 命令入口 ──────────────────────────────────────────

    def handle(self, args):
        """处理工具请求"""
        parts = args.strip().split(maxsplit=1)
        if not parts:
            return self.usage

        action = parts[0].lower()
        rest = parts[1] if len(parts) > 1 else ""

        if action == "help":
            return self.usage

        elif action == "connect":
            return self._do_connect(rest.strip())

        elif action == "disconnect":
            return self._do_disconnect()

        elif action == "qr":
            return self._do_qr()

        elif action == "status":
            return self._do_status()

        elif action == "strength":
            return self._do_strength(rest)

        elif action == "pulse":
            return self._do_pulse(rest)

        elif action == "wave":
            return self._do_wave(rest)

        elif action == "clear":
            return self._do_clear(rest)

        elif action == "limit":
            return self._do_limit(rest)

        else:
            return f"未知操作: {action}\n" + self.usage

    def _do_connect(self, url):
        """连接WebSocket服务器"""
        if not HAS_WS:
            return "❌ 缺少 websocket-client 库，请运行: pip install websocket-client"

        if self.connected:
            return f"已连接到 {self.ws_url}，请先 disconnect"

        if not url:
            return "请提供WebSocket地址，如: ws://192.168.1.100:9999"

        self.ws_url = url
        self.client_id = str(uuid.uuid4())
        connect_url = f"{url}/{self.client_id}"

        try:
            self.ws = websocket.WebSocketApp(
                connect_url,
                on_message=self._on_message,
                on_error=self._on_error,
                on_close=self._on_close,
                on_open=self._on_open
            )
            t = threading.Thread(target=self.ws.run_forever, daemon=True)
            t.start()

            for _ in range(30):
                if self.connected:
                    break
                time.sleep(0.1)

            if self.connected:
                return (
                    f"✅ 已连接到 {url}\n"
                    f"clientId: {self.client_id}\n"
                    f"请使用 /dglab qr 生成二维码供APP扫码配对"
                )
            else:
                return f"❌ 连接超时，请检查地址是否正确: {url}"
        except Exception as e:
            return f"❌ 连接失败: {e}"

    def _do_disconnect(self):
        """断开连接"""
        if self.ws:
            self._running = False
            try:
                self.ws.close()
            except Exception:
                pass
            self.connected = False
            self.paired = False
            return "已断开连接"
        return "当前没有连接"

    def _do_qr(self):
        """生成配对二维码"""
        if not self.connected or not self.client_id:
            return "❌ 请先连接到WebSocket服务器"

        qr_content = (
            f"https://www.dungeon-lab.com/app-download.php"
            f"#DGLAB-SOCKET#{self.ws_url}/{self.client_id}"
        )

        result = f"📱 配对二维码内容:\n{qr_content}\n\n"
        result += "请用 DG-LAB APP 扫描此二维码进行配对\n"

        if HAS_QRCODE:
            try:
                qr = qrcode.QRCode(version=1, box_size=10, border=4)
                qr.add_data(qr_content)
                qr.make(fit=True)
                qr.print_ascii(invert=True)
                result += "\n(已在终端显示ASCII二维码)"
            except Exception as e:
                result += f"\n(二维码生成失败: {e})"
        else:
            result += "\n💡 安装 qrcode 库可在终端显示ASCII二维码: pip install qrcode"

        return result

    def _do_status(self):
        """查看状态"""
        lines = ["📊 DG-Lab 连接状态:"]
        lines.append(f"  连接: {'✅ 已连接' if self.connected else '❌ 未连接'}")
        lines.append(f"  配对: {'✅ 已配对' if self.paired else '❌ 未配对'}")
        if self.connected:
            lines.append(f"  服务器: {self.ws_url}")
            lines.append(f"  clientId: {self.client_id}")
            if self.paired:
                lines.append(f"  targetId: {self.target_id}")
                lines.append(f"  通道A强度: {self.strength_a}/{self.limit_a}")
                lines.append(f"  通道B强度: {self.strength_b}/{self.limit_b}")
        return "\n".join(lines)

    def _do_strength(self, args):
        """控制强度: /dglab strength <ch> <op> <val>"""
        parts = args.strip().split()
        if len(parts) < 3:
            return "格式: strength <a/b> <+|-|=> <值>\n示例: strength a + 10"

        ch = parts[0].lower()
        op = parts[1]
        try:
            val = int(parts[2])
        except ValueError:
            return "强度值必须是整数"

        channel = 1 if ch == "a" else 2

        if op == "+":
            ok, msg = self._send_strength(channel, 1, val)
            return f"{'✅' if ok else '❌'} 通道{ch.upper()} 强度+{val}" if ok else f"❌ {msg}"
        elif op == "-":
            ok, msg = self._send_strength(channel, 0, val)
            return f"{'✅' if ok else '❌'} 通道{ch.upper()} 强度-{val}" if ok else f"❌ {msg}"
        elif op in ("=", "set"):
            val = max(0, min(200, val))
            ok, msg = self._send_strength(channel, 2, val)
            return f"{'✅' if ok else '❌'} 通道{ch.upper()} 强度设为 {val}" if ok else f"❌ {msg}"
        else:
            return f"未知操作: {op}，支持 +, -, ="

    def _do_pulse(self, args):
        """发送波形: /dglab pulse <ch> <freq> <intensity> <duration>"""
        parts = args.strip().split()
        if len(parts) < 3:
            return "格式: pulse <a/b> <频率10-240> <强度0-100> [持续秒数]"

        ch = parts[0].lower()
        try:
            freq = int(parts[1])
            intensity = int(parts[2])
            duration = int(parts[3]) if len(parts) > 3 else 5
        except ValueError:
            return "参数必须是整数"

        ok, msg = self._send_pulse(ch, freq, intensity, duration)
        return f"✅ {msg}" if ok else f"❌ {msg}"

    def _do_wave(self, args):
        """发送原始波形HEX: /dglab wave <ch> <hex>"""
        parts = args.strip().split(maxsplit=1)
        if len(parts) < 2:
            return "格式: wave <a/b> <HEX数据>\n示例: wave a 0A0A0A0A64646464"

        ch = parts[0].lower()
        hex_data = parts[1].strip()

        ok, msg = self._send_raw_wave(ch, hex_data)
        return f"✅ {msg}" if ok else f"❌ {msg}"

    def _do_clear(self, args):
        """清空波形队列"""
        ch = args.strip().lower()
        if ch in ("a", "1"):
            ok, msg = self._send_to_app("msg", "clear-1")
            return "✅ 已清空通道A波形队列" if ok else f"❌ {msg}"
        elif ch in ("b", "2"):
            ok, msg = self._send_to_app("msg", "clear-2")
            return "✅ 已清空通道B波形队列" if ok else f"❌ {msg}"
        elif ch == "both":
            self._send_to_app("msg", "clear-1")
            self._send_to_app("msg", "clear-2")
            return "✅ 已清空AB双通道波形队列"
        else:
            return "格式: clear <a/b/both>"

    def _do_limit(self, args):
        """设置软上限"""
        parts = args.strip().split()
        if len(parts) < 2:
            return "格式: limit <a/b> <0-200>"

        ch = parts[0].lower()
        try:
            val = max(0, min(200, int(parts[1])))
        except ValueError:
            return "上限值必须是0-200的整数"

        if ch in ("a", "1"):
            return (
                f"⚠️ 软上限需要通过蓝牙BF指令设置\n"
                f"  当前通道A上限: {self.limit_a}\n"
                f"  请在DG-LAB APP中调整软上限"
            )
        elif ch in ("b", "2"):
            return (
                f"⚠️ 软上限需要通过蓝牙BF指令设置\n"
                f"  当前通道B上限: {self.limit_b}\n"
                f"  请在DG-LAB APP中调整软上限"
            )
        else:
            return "格式: limit <a/b> <0-200>"