pythonTest/主播工具.py

import sys
import json
import threading
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QHBoxLayout, 
                             QVBoxLayout, QListWidget, QListWidgetItem, QCheckBox, 
                             QPushButton, QLineEdit, QTextEdit, QGroupBox, 
                             QSplitter, QLabel, QMessageBox, QComboBox, QShortcut)
from PyQt5.QtCore import Qt, QTimer, pyqtSignal, QObject
from PyQt5.QtGui import QKeySequence
import websocket
import time

class WebSocketClient(QObject):
    """WebSocket客户端，用于与服务器通信"""
    user_list_updated = pyqtSignal(list)  # 用户列表更新信号
    message_received = pyqtSignal(dict)   # 接收消息信号
    connection_status = pyqtSignal(bool)  # 连接状态信号
    
    def __init__(self, type_str="抖音自动评论-153"):
        super().__init__()
        self.type_str = type_str
        self.ws = None
        self.ws_user_info = None
        self.is_connected = False
        self.running = True
        self.is_admin = True
        self.user_id = None
        self.ws_thread = None
        
    def start(self):
        """启动WebSocket连接"""
        self.running = True
        self.ws_thread = threading.Thread(target=self._connect_websocket, daemon=True)
        self.ws_thread.start()
        
    def stop(self):
        """停止WebSocket连接"""
        self.running = False
        if self.ws:
            try:
                self.ws.close()
            except:
                pass
    
    def _connect_websocket(self):
        """WebSocket连接逻辑（在独立线程中运行）"""
        time_connect = 0
        
        def reconnect():
            nonlocal time_connect
            while self.running:
                try:
                    time_connect += 1
                    print(f"第{time_connect}次重连")
                    
                    # 创建WebSocket连接
                    self.ws = websocket.WebSocketApp(
                        'wss://ws.lamp.run',
                        on_open=self._on_open,
                        on_message=self._on_message,
                        on_error=self._on_error,
                        on_close=self._on_close
                    )
                    
                    # 运行WebSocket
                    self.ws.run_forever()
                    
                    # 如果连接断开且还在运行中，等待后重连
                    if self.running:
                        time.sleep(3)
                        
                except Exception as e:
                    print(f"WebSocket错误: {e}")
                    if self.running:
                        time.sleep(3)
        
        # 启动连接
        reconnect()
    
    def _on_open(self, ws):
        """WebSocket打开回调"""
        login_data = {
            "route": "login",
            "type": self.type_str,
            "admin": self.is_admin,
            "id": "PYQT_CLIENT"
        }
        ws.send(json.dumps(login_data))
        self.is_connected = True
        self.connection_status.emit(True)
        print("WebSocket连接成功")
        
        # 延迟一下，等登录完成后再获取列表
        def get_list():
            time.sleep(0.5)
            if self.ws and self.is_connected:
                self.send("getList")
        
        threading.Thread(target=get_list, daemon=True).start()
    
    def _on_message(self, ws, message):
        """接收消息回调"""
        try:
            data = json.loads(message)
            
            if data.get('type') == 'userInfo':
                self.ws_user_info = data.get('value')
                if self.ws_user_info:
                    self.user_id = self.ws_user_info.get('userID')
                print(f"用户信息: {data.get('value')}")
                
            elif data.get('type') == 'getList':
                # 接收到用户列表
                user_list = data.get('value', [])
                # 提取用户ID列表
                user_ids = []
                for user in user_list:
                    user_id = user.get('id')
                    if user_id and user_id != 'PYQT_CLIENT':  # 排除自己
                        user_ids.append(user_id)
                
                # 发送信号更新UI
                self.user_list_updated.emit(user_ids)
                print(f"收到用户列表: {user_ids}")
                
            else:
                # 其他消息
                self.message_received.emit(data)
                
        except json.JSONDecodeError as e:
            print(f"JSON解析错误: {e}")
        except Exception as e:
            print(f"消息处理错误: {e}")
    
    def _on_error(self, ws, error):
        """WebSocket错误回调"""
        print(f"WebSocket错误: {error}")
        self.is_connected = False
        self.connection_status.emit(False)
    
    def _on_close(self, ws, close_status_code, close_msg):
        """WebSocket关闭回调"""
        print("WebSocket连接关闭")
        self.is_connected = False
        self.connection_status.emit(False)
    
    def send(self, route, value=None, send_to=None):
        """发送消息
        
        Args:
            route: 路由名称，如 "sendmessage"
            value: 要发送的内容（话术）
            send_to: 接收方的userID
        """
        if not self.ws or not self.is_connected:
            print("WebSocket未连接")
            return False
        
        try:
            message = {
                "route": route,
                "type": self.type_str,
            }
            
            # 添加userID（发送者的ID）
            if self.user_id:
                message["userID"] = self.user_id
            
            # 添加接收者ID
            if send_to:
                message["id"] = "抖音自动评论-153-" + send_to
            
            # 添加内容
            if value is not None:
                message["value"] = value
            
            # 发送消息
            self.ws.send(json.dumps(message))
            print(f"发送消息: {message}")
            return True
            
        except Exception as e:
            print(f"发送消息失败: {e}")
            return False
    
    def get_user_list(self):
        """主动获取用户列表"""
        if self.is_connected:
            return self.send("getList")
        return False


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("三栏布局 - 用户独立话术管理 (快捷键F1-F5)")
        self.setGeometry(100, 100, 1300, 750)
        
        # 数据存储：用户话术字典 {用户名: [话术列表]}
        self.user_scripts = {}
        # 当前选中的用户
        self.current_user = None
        # WebSocket客户端
        self.ws_client = None
        # 用户列表缓存
        self.user_cache = []
        
        # 中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 主布局 - 水平分割器
        main_splitter = QSplitter(Qt.Horizontal)
        central_layout = QHBoxLayout(central_widget)
        central_layout.addWidget(main_splitter)
        central_layout.setContentsMargins(10, 10, 10, 10)
        
        # ---------- 左侧：日志显示 ----------
        left_widget = QWidget()
        left_layout = QVBoxLayout(left_widget)
        left_layout.setContentsMargins(5, 5, 5, 5)
        
        log_label = QLabel("📋 运行日志")
        log_label.setStyleSheet("font-weight: bold; font-size: 14px;")
        left_layout.addWidget(log_label)
        
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setPlaceholderText("日志信息将显示在这里...")
        self.log_text.append("[系统] 程序启动")
        self.log_text.append("[系统] 界面初始化完成")
        self.log_text.append("[系统] 快捷键: F1-F5 快速切换用户")
        left_layout.addWidget(self.log_text)
        
        log_btn_layout = QHBoxLayout()
        clear_log_btn = QPushButton("清空日志")
        clear_log_btn.clicked.connect(self.clear_log)
        log_btn_layout.addStretch()
        log_btn_layout.addWidget(clear_log_btn)
        left_layout.addLayout(log_btn_layout)
        
        # ---------- 中间：在线用户列表 ----------
        middle_widget = QWidget()
        middle_layout = QVBoxLayout(middle_widget)
        middle_layout.setContentsMargins(5, 5, 5, 5)
        
        user_label = QLabel("👥 在线用户列表")
        user_label.setStyleSheet("font-weight: bold; font-size: 14px;")
        middle_layout.addWidget(user_label)
        
        # 连接状态和控制按钮
        control_layout = QHBoxLayout()
        self.status_label = QLabel("⚪ 未连接")
        self.status_label.setStyleSheet("color: red;")
        control_layout.addWidget(self.status_label)
        control_layout.addStretch()
        
        self.refresh_btn = QPushButton("🔄 刷新列表")
        self.refresh_btn.clicked.connect(self.refresh_user_list)
        self.refresh_btn.setEnabled(False)
        control_layout.addWidget(self.refresh_btn)
        middle_layout.addLayout(control_layout)
        
        select_btn_layout = QHBoxLayout()
        self.select_all_btn = QPushButton("✅ 全选")
        self.deselect_all_btn = QPushButton("❌ 取消全选")
        self.select_all_btn.clicked.connect(self.select_all_users)
        self.deselect_all_btn.clicked.connect(self.deselect_all_users)
        select_btn_layout.addWidget(self.select_all_btn)
        select_btn_layout.addWidget(self.deselect_all_btn)
        select_btn_layout.addStretch()
        middle_layout.addLayout(select_btn_layout)
        
        self.user_list_widget = QListWidget()
        self.user_list_widget.setSelectionMode(QListWidget.NoSelection)
        middle_layout.addWidget(self.user_list_widget)
        
        self.user_count_label = QLabel("当前在线人数: 0")
        self.user_count_label.setStyleSheet("color: gray;")
        middle_layout.addWidget(self.user_count_label)
        
        # ---------- 右侧：用户独立话术列表 ----------
        right_widget = QWidget()
        right_layout = QVBoxLayout(right_widget)
        right_layout.setContentsMargins(5, 5, 5, 5)
        
        # 顶部：用户选择下拉框
        user_select_layout = QHBoxLayout()
        user_select_layout.addWidget(QLabel("选择用户 (快捷键F1-F5):"))
        self.user_combo = QComboBox()
        self.user_combo.currentTextChanged.connect(self.on_user_changed)
        user_select_layout.addWidget(self.user_combo)
        user_select_layout.addStretch()
        right_layout.addLayout(user_select_layout)
        
        # 话术列表标题
        script_label = QLabel("📝 话术列表 (当前用户)")
        script_label.setStyleSheet("font-weight: bold; font-size: 14px;")
        right_layout.addWidget(script_label)
        
        # 话术列表显示
        self.script_list_widget = QListWidget()
        right_layout.addWidget(self.script_list_widget)
        
        # 话术编辑区域
        script_edit_layout = QVBoxLayout()
        
        # 添加话术
        add_layout = QHBoxLayout()
        self.script_input = QLineEdit()
        self.script_input.setPlaceholderText("输入新话术内容...")
        add_btn = QPushButton("➕ 添加")
        add_btn.clicked.connect(self.add_script)
        add_layout.addWidget(self.script_input)
        add_layout.addWidget(add_btn)
        script_edit_layout.addLayout(add_layout)
        
        # 删除和修改按钮
        modify_layout = QHBoxLayout()
        delete_btn = QPushButton("🗑️ 删除选中")
        delete_btn.clicked.connect(self.delete_script)
        modify_btn = QPushButton("✏️ 修改选中")
        modify_btn.clicked.connect(self.modify_script)
        modify_layout.addWidget(delete_btn)
        modify_layout.addWidget(modify_btn)
        script_edit_layout.addLayout(modify_layout)
        
        # 保存按钮
        save_layout = QHBoxLayout()
        save_layout.addStretch()
        self.save_scripts_btn = QPushButton("💾 保存话术配置")
        self.save_scripts_btn.setStyleSheet("background-color: #2196F3; color: white; font-weight: bold; padding: 8px 20px; border-radius: 5px;")
        self.save_scripts_btn.clicked.connect(self.save_user_scripts)
        save_layout.addWidget(self.save_scripts_btn)
        script_edit_layout.addLayout(save_layout)
        
        right_layout.addLayout(script_edit_layout)
        
        # 发送任务按钮
        send_layout = QHBoxLayout()
        send_layout.addStretch()
        self.send_task_btn = QPushButton("🚀 发送任务")
        self.send_task_btn.setStyleSheet("background-color: #4CAF50; color: white; font-weight: bold; padding: 8px 20px; border-radius: 5px;")
        self.send_task_btn.clicked.connect(self.send_task)
        send_layout.addWidget(self.send_task_btn)
        right_layout.addLayout(send_layout)
        
        # ---------- 将三个部件加入分割器 ----------
        main_splitter.addWidget(left_widget)
        main_splitter.addWidget(middle_widget)
        main_splitter.addWidget(right_widget)
        main_splitter.setSizes([300, 450, 450])
        
        # ---------- 设置快捷键 F1-F5 ----------
        self.setup_shortcuts()
        
        # 窗口样式
        self.setStyleSheet("""
            QGroupBox {
                font-weight: bold;
                border: 1px solid #ccc;
                border-radius: 5px;
                margin-top: 10px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 5px 0 5px;
            }
            QListWidget::item {
                padding: 5px;
            }
            QTextEdit {
                background-color: #f8f8f8;
                font-family: Consolas, monospace;
            }
            QComboBox {
                padding: 5px;
                min-width: 150px;
            }
        """)
        
        # 连接WebSocket
        self.connect_websocket()
    
    # ---------- WebSocket连接管理 ----------
    def connect_websocket(self):
        """连接到WebSocket服务器"""
        self.log_text.append("[系统] 正在连接WebSocket服务器...")
        self.status_label.setText("🔄 正在连接...")
        self.status_label.setStyleSheet("color: orange;")
        
        # 创建WebSocket客户端
        self.ws_client = WebSocketClient("抖音自动评论-153")
        
        # 连接信号
        self.ws_client.user_list_updated.connect(self.on_user_list_updated)
        self.ws_client.message_received.connect(self.on_ws_message)
        self.ws_client.connection_status.connect(self.on_connection_status)
        
        # 启动WebSocket
        self.ws_client.start()
        
        # 定时器检查连接状态
        self.status_timer = QTimer()
        self.status_timer.timeout.connect(self.check_connection)
        self.status_timer.start(5000)  # 每5秒检查一次
    
    def on_connection_status(self, connected):
        """WebSocket连接状态变化"""
        if connected:
            self.status_label.setText("🟢 已连接")
            self.status_label.setStyleSheet("color: green;")
            self.refresh_btn.setEnabled(True)
            self.log_text.append("[系统] WebSocket连接成功")
            # 延迟获取用户列表
            QTimer.singleShot(1000, self.refresh_user_list)
        else:
            self.status_label.setText("🔴 已断开")
            self.status_label.setStyleSheet("color: red;")
            self.refresh_btn.setEnabled(False)
            self.log_text.append("[系统] WebSocket连接断开")
    
    def check_connection(self):
        """检查WebSocket连接状态"""
        if self.ws_client:
            if not self.ws_client.is_connected:
                # 尝试重新获取列表
                self.ws_client.get_user_list()
    
    def refresh_user_list(self):
        """刷新用户列表"""
        if self.ws_client and self.ws_client.is_connected:
            self.log_text.append("[系统] 刷新用户列表...")
            self.ws_client.get_user_list()
        else:
            self.log_text.append("[系统] WebSocket未连接，无法刷新列表")
    
    def on_user_list_updated(self, user_list):
        """用户列表更新"""
        self.user_cache = user_list
        self.update_user_list(user_list)
        self.log_text.append(f"[系统] 更新用户列表，当前在线: {len(user_list)}人")
    
    def on_ws_message(self, message):
        """WebSocket消息接收"""
        # 处理其他类型的消息
        msg_type = message.get('type')
        if msg_type:
            self.log_text.append(f"[WebSocket] 收到消息: {msg_type}")
    
    # ---------- 用户管理 ----------
    def update_user_list(self, user_list):
        """更新用户列表UI"""
        # 保存当前选中的用户
        selected_users = self.get_selected_users()
        current_combo_text = self.user_combo.currentText()
        
        # 清空列表
        self.user_list_widget.clear()
        
        # 添加用户
        for user_id in user_list:
            self.add_user_item(user_id)
            # 如果用户不存在，初始化话术数据
            if user_id not in self.user_scripts:
                self.user_scripts[user_id] = [f"{user_id}，666"]
        
        # 更新下拉框
        self.user_combo.clear()
        self.user_combo.addItems(user_list)
        
        # 恢复选中状态
        if current_combo_text in user_list:
            self.user_combo.setCurrentText(current_combo_text)
        elif user_list:
            self.user_combo.setCurrentIndex(0)
        
        # 恢复复选框选中状态
        for user in selected_users:
            if user in user_list:
                for i in range(self.user_list_widget.count()):
                    item = self.user_list_widget.item(i)
                    widget = self.user_list_widget.itemWidget(item)
                    if isinstance(widget, QCheckBox) and widget.text() == user:
                        widget.setChecked(True)
                        break
        
        # 更新人数
        self.user_count_label.setText(f"当前在线人数: {len(user_list)}")
    
    def add_user_item(self, username):
        """添加一个带复选框的用户项"""
        item = QListWidgetItem()
        checkbox = QCheckBox(username)
        self.user_list_widget.addItem(item)
        self.user_list_widget.setItemWidget(item, checkbox)
        checkbox.stateChanged.connect(self.on_user_check_changed)
    
    def get_user_checkboxes(self):
        """获取所有用户复选框"""
        checkboxes = []
        for i in range(self.user_list_widget.count()):
            item = self.user_list_widget.item(i)
            widget = self.user_list_widget.itemWidget(item)
            if isinstance(widget, QCheckBox):
                checkboxes.append(widget)
        return checkboxes
    
    def get_selected_users(self):
        """获取所有选中的用户"""
        selected = []
        for cb in self.get_user_checkboxes():
            if cb.isChecked():
                selected.append(cb.text())
        return selected
    
    def select_all_users(self):
        for cb in self.get_user_checkboxes():
            cb.setChecked(True)
        self.log_text.append("[操作] 已全选所有用户")
    
    def deselect_all_users(self):
        for cb in self.get_user_checkboxes():
            cb.setChecked(False)
        self.log_text.append("[操作] 已取消全选所有用户")
    
    def on_user_check_changed(self):
        """复选框状态改变时更新下拉框列表"""
        # 获取当前所有用户
        current_users = []
        for i in range(self.user_list_widget.count()):
            item = self.user_list_widget.item(i)
            widget = self.user_list_widget.itemWidget(item)
            if isinstance(widget, QCheckBox):
                current_users.append(widget.text())
        
        # 更新下拉框
        current_selection = self.user_combo.currentText()
        self.user_combo.clear()
        self.user_combo.addItems(current_users)
        
        # 如果之前的用户还在列表中，恢复选择
        if current_selection in current_users:
            self.user_combo.setCurrentText(current_selection)
        elif current_users:
            self.user_combo.setCurrentIndex(0)
    
    # ---------- 快捷键设置 ----------
    def setup_shortcuts(self):
        """设置F1-F5快捷键切换用户"""
        shortcuts = [
            (Qt.Key_F1, 0),
            (Qt.Key_F2, 1),
            (Qt.Key_F3, 2),
            (Qt.Key_F4, 3),
            (Qt.Key_F5, 4),
        ]
        
        for key, index in shortcuts:
            shortcut = QShortcut(QKeySequence(key), self)
            shortcut.activated.connect(lambda idx=index: self.switch_to_user_by_index(idx))
    
    def switch_to_user_by_index(self, index):
        """根据索引切换到对应用户"""
        if index < self.user_combo.count():
            user = self.user_combo.itemText(index)
            if user:
                self.user_combo.setCurrentText(user)
                self.log_text.append(f"[快捷键] 切换到用户: {user}")
        else:
            self.log_text.append(f"[快捷键] 用户索引 {index} 不存在")
    
    # ---------- 话术管理 (用户独立) ----------
    def load_user_scripts(self, username):
        """加载指定用户的话术列表"""
        if username in self.user_scripts:
            self.script_list_widget.clear()
            for script in self.user_scripts[username]:
                self.script_list_widget.addItem(script)
            self.current_user = username
            self.log_text.append(f"[系统] 加载用户 '{username}' 的话术列表")
    
    def on_user_changed(self, username):
        """下拉框切换用户"""
        if username and username != self.current_user:
            # 如果当前用户有修改但未保存，先保存
            if self.current_user and self.current_user in self.user_scripts:
                self.save_current_user_scripts()
            
            self.load_user_scripts(username)
    
    def save_current_user_scripts(self):
        """保存当前用户的话术到内存"""
        if self.current_user and self.current_user in self.user_scripts:
            scripts = []
            for i in range(self.script_list_widget.count()):
                scripts.append(self.script_list_widget.item(i).text())
            self.user_scripts[self.current_user] = scripts
            self.log_text.append(f"[系统] 自动保存用户 '{self.current_user}' 的话术配置")
    
    def save_user_scripts(self):
        """手动保存当前用户的话术（并持久化）"""
        if not self.current_user:
            QMessageBox.warning(self, "提示", "请先选择一个用户！")
            return
        
        # 保存到内存
        self.save_current_user_scripts()
        
        # 持久化保存到文件
        try:
            with open("user_scripts_backup.json", "w", encoding="utf-8") as f:
                json.dump(self.user_scripts, f, ensure_ascii=False, indent=2)
            self.log_text.append(f"[系统] 话术配置已保存到文件 user_scripts_backup.json")
            QMessageBox.information(self, "保存成功", f"用户 '{self.current_user}' 的话术配置已保存！")
        except Exception as e:
            QMessageBox.critical(self, "保存失败", f"保存到文件时出错: {str(e)}")
    
    def add_script(self):
        """添加话术到当前用户"""
        text = self.script_input.text().strip()
        if text:
            self.script_list_widget.addItem(text)
            self.script_input.clear()
            self.save_current_user_scripts()
            self.log_text.append(f"[操作] 为用户 '{self.current_user}' 添加话术: {text}")
        else:
            QMessageBox.information(self, "提示", "话术内容不能为空！")
    
    def delete_script(self):
        """删除当前用户选中的话术"""
        current_row = self.script_list_widget.currentRow()
        if current_row >= 0:
            item_text = self.script_list_widget.item(current_row).text()
            self.script_list_widget.takeItem(current_row)
            self.save_current_user_scripts()
            self.log_text.append(f"[操作] 删除用户 '{self.current_user}' 的话术: {item_text}")
        else:
            QMessageBox.information(self, "提示", "请先选中一条话术！")
    
    def modify_script(self):
        """修改当前用户选中的话术"""
        current_row = self.script_list_widget.currentRow()
        if current_row >= 0:
            old_text = self.script_list_widget.item(current_row).text()
            new_text, ok = QMessageBox.getText(self, "修改话术", "请输入新的话术内容:", 
                                               text=old_text)
            if ok and new_text.strip():
                self.script_list_widget.item(current_row).setText(new_text.strip())
                self.save_current_user_scripts()
                self.log_text.append(f"[操作] 修改用户 '{self.current_user}' 的话术: {old_text} -> {new_text.strip()}")
            elif ok and not new_text.strip():
                QMessageBox.information(self, "提示", "话术内容不能为空！")
        else:
            QMessageBox.information(self, "提示", "请先选中一条话术！")
    
    def clear_log(self):
        self.log_text.clear()
        self.log_text.append("[系统] 日志已清空")
        self.log_text.append("[系统] 快捷键: F1-F5 快速切换用户")
    
    # ---------- 发送任务 ----------
    def send_task(self):
        """发送任务到选中的用户"""
        selected_users = self.get_selected_users()
        
        if not selected_users:
            QMessageBox.warning(self, "提示", "请至少选择一个用户！")
            return
        
        # 获取当前选中的话术
        current_row = self.script_list_widget.currentRow()
        if current_row < 0:
            QMessageBox.warning(self, "提示", "请先在话术列表中选中一条话术！")
            return
        
        selected_script = self.script_list_widget.item(current_row).text()
        
        # 记录日志
        self.log_text.append(f"[任务] 发送话术: {selected_script}")
        self.log_text.append(f"[任务] 发送给 {len(selected_users)} 个用户:")
        
        # 通过WebSocket发送任务
        if self.ws_client and self.ws_client.is_connected:
            success_count = 0
            for user in selected_users:
                # 发送给指定用户，路由为 "sendmessage"
                if self.ws_client.send("sendmessage", selected_script, user):
                    success_count += 1
                    self.log_text.append(f"[任务] 已发送给用户: {user}")
                else:
                    self.log_text.append(f"[任务] 发送给用户 {user} 失败")
            
            self.log_text.append(f"[任务] 发送完成，成功发送给 {success_count}/{len(selected_users)} 个用户")
        else:
            self.log_text.append("[任务] WebSocket未连接，无法发送")
    
    # ---------- 窗口关闭时保存 ----------
    def closeEvent(self, event):
        """关闭窗口时自动保存所有用户话术"""
        if self.current_user:
            self.save_current_user_scripts()
        
        try:
            with open("user_scripts_backup.json", "w", encoding="utf-8") as f:
                json.dump(self.user_scripts, f, ensure_ascii=False, indent=2)
            self.log_text.append("[系统] 程序退出，已自动保存所有话术配置")
        except:
            pass
        
        # 关闭WebSocket连接
        if self.ws_client:
            self.ws_client.stop()
        
        event.accept()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())