import subprocess
import psutil
import os
import sys
import json
import socket
import time
import re
from loguru import logger
from PySide6.QtWidgets import (QApplication, QMainWindow, QPushButton, QMessageBox,
                               QVBoxLayout, QWidget, QLabel, QGroupBox, QGridLayout,
                               QStatusBar, QFrame, QDialog, QLineEdit, QFormLayout,
                               QDialogButtonBox, QFileDialog, QHBoxLayout, QComboBox,
                               QSpinBox, QCheckBox)
from PySide6.QtCore import Qt, QThread, Signal
from PySide6.QtGui import QFont

logger.add(sys.stderr, format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}", level="INFO")

CONFIG_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "dual_screen_config.json")

def load_config():
    if os.path.exists(CONFIG_FILE):
        with open(CONFIG_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    return {
        "server_ip": "192.168.1.100",
        "mediamtx_path": r"D:\ScreenCast\mediamtx\mediamtx.exe",
        "ffmpeg_path": r"D:\ScreenCast\ffmpeg\bin\ffmpeg.exe",
        "front_stream_path": "front",
        "back_stream_path": "back"
    }

def save_config(config):
    with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
        json.dump(config, f, ensure_ascii=False, indent=4)
    logger.info(f"配置已保存到 {CONFIG_FILE}")

def get_displays():
    displays = []
    try:
        result = subprocess.run(
            ['powershell', '-Command',
             'Add-Type -AssemblyName System.Windows.Forms; '
             '[System.Windows.Forms.Screen]::AllScreens | ForEach-Object { '
             '$dev = $_.DeviceName; '
             '$w = $_.Bounds.Width; '
             '$h = $_.Bounds.Height; '
             '$x = $_.Bounds.X; '
             '$y = $_.Bounds.Y; '
             '$isPrimary = $_.Primary; '
             'Write-Output "$dev|$w|$h|$x|$y|$isPrimary" }'],
            capture_output=True, text=True, timeout=10
        )
        for line in result.stdout.strip().split('\n'):
            if line:
                parts = line.split('|')
                if len(parts) >= 6:
                    displays.append({
                        'name': parts[0],
                        'width': int(parts[1]),
                        'height': int(parts[2]),
                        'x': int(parts[3]),
                        'y': int(parts[4]),
                        'primary': parts[5] == 'True'
                    })
    except Exception as e:
        logger.error(f"获取显示器信息失败: {e}")
    return displays

class ConfigDialog(QDialog):
    def __init__(self, config, parent=None):
        super().__init__(parent)
        self.config = config
        self.setWindowTitle("设置")
        self.setFixedSize(450, 280)

        layout = QFormLayout(self)

        self.server_ip_edit = QLineEdit(config["server_ip"])
        self.mediamtx_path_edit = QLineEdit(config["mediamtx_path"])
        self.ffmpeg_path_edit = QLineEdit(config["ffmpeg_path"])
        self.front_stream_edit = QLineEdit(config["front_stream_path"])
        self.back_stream_edit = QLineEdit(config["back_stream_path"])

        self.mediamtx_btn = QPushButton("浏览...")
        self.mediamtx_btn.clicked.connect(lambda: self.browse("mediamtx"))
        self.ffmpeg_btn = QPushButton("浏览...")
        self.ffmpeg_btn.clicked.connect(lambda: self.browse("ffmpeg"))

        path_layout = QGridLayout()
        path_layout.addWidget(self.mediamtx_path_edit, 0, 0)
        path_layout.addWidget(self.mediamtx_btn, 0, 1)
        path_layout.addWidget(self.ffmpeg_path_edit, 1, 0)
        path_layout.addWidget(self.ffmpeg_btn, 1, 1)

        layout.addRow("服务器IP:", self.server_ip_edit)
        layout.addRow("MediaMTX路径:", self.mediamtx_path_edit)
        layout.addRow("FFmpeg路径:", self.ffmpeg_path_edit)
        layout.addRow("前置屏幕流名:", self.front_stream_edit)
        layout.addRow("后置屏幕流名:", self.back_stream_edit)

        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        layout.addRow(buttons)

    def browse(self, field):
        if field == "mediamtx":
            path = QFileDialog.getOpenFileName(self, "选择MediaMTX", "", "Executable (*.exe)")[0]
            if path:
                self.mediamtx_path_edit.setText(path)
        elif field == "ffmpeg":
            path = QFileDialog.getOpenFileName(self, "选择FFmpeg", "", "Executable (*.exe)")[0]
            if path:
                self.ffmpeg_path_edit.setText(path)

    def get_config(self):
        return {
            "server_ip": self.server_ip_edit.text().strip(),
            "mediamtx_path": self.mediamtx_path_edit.text().strip(),
            "ffmpeg_path": self.ffmpeg_path_edit.text().strip(),
            "front_stream_path": self.front_stream_edit.text().strip() or "front",
            "back_stream_path": self.back_stream_edit.text().strip() or "back"
        }

class ConnectionChecker(QThread):
    status_update = Signal(dict)

    def run(self):
        while True:
            cfg = self.config
            status = {
                "mediamtx": self.check_mediamtx(cfg["mediamtx_path"]),
                "ffmpeg": self.check_ffmpeg(cfg["ffmpeg_path"]),
                "server": self.check_server(cfg["server_ip"]),
                "server_port": self.check_port(cfg["server_ip"], 8554),
            }
            self.status_update.emit(status)
            time.sleep(3)

    def set_config(self, config):
        self.config = config

    def check_mediamtx(self, path):
        for proc in psutil.process_iter(['name']):
            if proc.info['name'] == 'mediamtx.exe':
                return True
        return False

    def check_ffmpeg(self, path):
        return os.path.exists(path)

    def check_server(self, ip):
        try:
            socket.create_connection((ip, 8554), timeout=2)
            return True
        except:
            return False

    def check_port(self, ip, port):
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(2)
            result = sock.connect_ex((ip, port))
            sock.close()
            return result == 0
        except:
            return False

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.config = load_config()
        self.processes = []
        self.setWindowTitle("双屏投屏控制")
        self.setFixedSize(600, 500)

        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)

        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        title_layout = QHBoxLayout()
        title_layout.addStretch()
        self.btn_settings = QPushButton("⚙ 设置")
        self.btn_settings.setFixedSize(70, 30)
        self.btn_settings.clicked.connect(self.open_settings)
        title_layout.addWidget(self.btn_settings)

        title_label = QLabel("双屏投屏系统")
        title_font = QFont("Microsoft YaHei", 14, QFont.Bold)
        title_label.setFont(title_font)
        title_label.setAlignment(Qt.AlignCenter)

        title_main_layout = QVBoxLayout()
        title_main_layout.addLayout(title_layout)
        title_main_layout.addWidget(title_label)
        main_layout.addLayout(title_main_layout)

        self.displays = get_displays()
        display_info = QLabel(f"检测到 {len(self.displays)} 个显示器")
        display_info.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(display_info)

        self.display_group = QGroupBox("显示器信息")
        display_layout = QGridLayout()

        self.display_labels = []
        for i, display in enumerate(self.displays):
            label = QLabel(f"显示器{i+1}: {display['width']}x{display['height']} "
                          f"(@ {display['x']},{display['y']}) "
                          f"{'主屏' if display['primary'] else ''}")
            label.setStyleSheet("padding: 5px;")
            display_layout.addWidget(label, i, 0)
            self.display_labels.append(label)

        self.display_group.setLayout(display_layout)
        main_layout.addWidget(self.display_group)

        config_group = QGroupBox("本地配置")
        config_layout = QGridLayout()

        self.mediamtx_label = QLabel("MediaMTX: 未检测")
        self.ffmpeg_label = QLabel("FFmpeg: 未检测")
        self.server_label = QLabel("服务器连接: 未检测")

        self.mediamtx_indicator = QLabel("●")
        self.mediamtx_indicator.setFixedWidth(20)
        self.ffmpeg_indicator = QLabel("●")
        self.ffmpeg_indicator.setFixedWidth(20)
        self.server_indicator = QLabel("●")
        self.server_indicator.setFixedWidth(20)

        config_layout.addWidget(self.mediamtx_indicator, 0, 0)
        config_layout.addWidget(self.mediamtx_label, 0, 1)
        config_layout.addWidget(self.ffmpeg_indicator, 1, 0)
        config_layout.addWidget(self.ffmpeg_label, 1, 1)
        config_layout.addWidget(self.server_indicator, 2, 0)
        config_layout.addWidget(self.server_label, 2, 1)

        config_group.setLayout(config_layout)
        main_layout.addWidget(config_group)

        self.info_group = QGroupBox("流信息")
        self.update_info_group()
        main_layout.addWidget(self.info_group)

        btn_layout = QHBoxLayout()

        self.btn_front = QPushButton("🎬 投送前置屏幕")
        self.btn_front.setFixedHeight(45)
        self.btn_front.setFont(QFont("Microsoft YaHei", 10))
        self.btn_front.setStyleSheet("""
            QPushButton {
                background-color: #1976D2;
                color: white;
                border: none;
                border-radius: 8px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #1565C0;
            }
        """)
        self.btn_front.clicked.connect(lambda: self.push_screen("front"))

        self.btn_back = QPushButton("🎬 投送后置屏幕")
        self.btn_back.setFixedHeight(45)
        self.btn_back.setFont(QFont("Microsoft YaHei", 10))
        self.btn_back.setStyleSheet("""
            QPushButton {
                background-color: #388E3C;
                color: white;
                border: none;
                border-radius: 8px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #2E7D32;
            }
        """)
        self.btn_back.clicked.connect(lambda: self.push_screen("back"))

        btn_layout.addWidget(self.btn_front)
        btn_layout.addWidget(self.btn_back)
        main_layout.addLayout(btn_layout)

        self.btn_stop = QPushButton("⏹ 停止所有推流")
        self.btn_stop.setFixedHeight(40)
        self.btn_stop.setFont(QFont("Microsoft YaHei", 10))
        self.btn_stop.setStyleSheet("""
            QPushButton {
                background-color: #D32F2F;
                color: white;
                border: none;
                border-radius: 8px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #C62828;
            }
        """)
        self.btn_stop.clicked.connect(self.stop_all)
        main_layout.addWidget(self.btn_stop)

        self.log_label = QLabel("日志: 等待操作...")
        self.log_label.setWordWrap(True)
        self.log_label.setFrameShape(QFrame.Shape.StyledPanel)
        self.log_label.setFixedHeight(50)
        main_layout.addWidget(self.log_label)

        self.checker = ConnectionChecker()
        self.checker.set_config(self.config)
        self.checker.status_update.connect(self.update_status)
        self.checker.start()

        logger.info("双屏投屏控制界面已初始化")

    def update_info_group(self):
        cfg = self.config
        if self.info_group.layout():
            QWidget().setLayout(self.info_group.layout())
        info_layout = QGridLayout()
        info_layout.addWidget(QLabel("服务器IP:"), 0, 0)
        info_layout.addWidget(QLabel(cfg["server_ip"]), 0, 1)
        info_layout.addWidget(QLabel("前置屏幕流:"), 1, 0)
        info_layout.addWidget(QLabel(f"rtsp://{cfg['server_ip']}:8554/{cfg['front_stream_path']}"), 1, 1)
        info_layout.addWidget(QLabel("后置屏幕流:"), 2, 0)
        info_layout.addWidget(QLabel(f"rtsp://{cfg['server_ip']}:8554/{cfg['back_stream_path']}"), 2, 1)
        self.info_group.setLayout(info_layout)

    def open_settings(self):
        dialog = ConfigDialog(self.config, self)
        if dialog.exec():
            self.config = dialog.get_config()
            save_config(self.config)
            self.checker.set_config(self.config)
            self.update_info_group()
            self.log_label.setText("日志: 配置已更新")

    def update_status(self, status):
        if status["mediamtx"]:
            self.mediamtx_indicator.setStyleSheet("color: green; font-size: 16px;")
            self.mediamtx_label.setText("MediaMTX: ✓ 运行中")
        else:
            self.mediamtx_indicator.setStyleSheet("color: red; font-size: 16px;")
            self.mediamtx_label.setText("MediaMTX: ✗ 未运行")

        if status["ffmpeg"]:
            self.ffmpeg_indicator.setStyleSheet("color: green; font-size: 16px;")
            self.ffmpeg_label.setText("FFmpeg: ✓ 已安装")
        else:
            self.ffmpeg_indicator.setStyleSheet("color: red; font-size: 16px;")
            self.ffmpeg_label.setText("FFmpeg: ✗ 未找到")

        if status["server"]:
            self.server_indicator.setStyleSheet("color: green; font-size: 16px;")
            self.server_label.setText("服务器连接: ✓ 可连接")
        else:
            self.server_indicator.setStyleSheet("color: orange; font-size: 16px;")
            self.server_label.setText("服务器连接: ✗ 连接失败")

    def push_screen(self, screen_type):
        cfg = self.config
        displays = self.displays

        if len(displays) < 2:
            QMessageBox.warning(self, "错误", "未检测到足够的显示器！")
            return

        if screen_type == "front":
            display = displays[0]
            stream_path = cfg["front_stream_path"]
            screen_name = "前置屏幕"
        else:
            display = displays[1] if len(displays) > 1 else displays[0]
            stream_path = cfg["back_stream_path"]
            screen_name = "后置屏幕"

        logger.info(f"开始投送{screen_name}: {display['width']}x{display['height']}")

        cmd = [
            cfg["ffmpeg_path"],
            "-f", "gdigrab",
            "-offset_x", str(display['x']),
            "-offset_y", str(display['y']),
            "-video_size", f"{display['width']}x{display['height']}",
            "-framerate", "30",
            "-i", "desktop",
            "-c:v", "libx264",
            "-preset", "ultrafast",
            "-tune", "zerolatency",
            "-f", "rtsp",
            f"rtsp://{cfg['server_ip']}:8554/{stream_path}"
        ]

        logger.info(f"执行FFmpeg命令: {' '.join(cmd)}")
        process = subprocess.Popen(cmd, creationflags=subprocess.CREATE_NEW_CONSOLE)
        self.processes.append(process)

        self.log_label.setText(f"日志: {screen_name}正在投送... "
                              f"访问 http://{cfg['server_ip']}:8889/{stream_path}")
        QMessageBox.information(self, "提示",
                              f"{screen_name}投送已启动！\n\n"
                              f"观众访问：http://{cfg['server_ip']}:8889/{stream_path}")

    def stop_all(self):
        for proc in self.processes:
            try:
                proc.terminate()
            except:
                pass
        self.processes.clear()

        for proc in psutil.process_iter(['name']):
            try:
                if proc.info['name'] == 'ffmpeg.exe':
                    proc.terminate()
            except:
                pass

        self.log_label.setText("日志: 所有推流已停止")
        logger.info("所有推流已停止")

    def closeEvent(self, event):
        self.stop_all()
        self.checker.terminate()
        event.accept()

if __name__ == "__main__":
    logger.info("应用程序启动")
    app = QApplication(sys.argv)
    app.setStyle("Fusion")
    window = MainWindow()
    window.show()
    sys.exit(app.exec())