gui.py

import sys
import os
import json
from threading import Event
from PyQt6.QtWidgets import *
from PyQt6.QtCore import *
from PyQt6.QtGui import *

# 导入原有模块
from login import BiliLogin
from parser import BiliParser
from utils import parse_url_input, get_season_info, get_fav_list
from config import DEBUG

# ---------- 轮询线程（扫码登录） ----------
class PollThread(QThread):
    status_signal = pyqtSignal(str)
    finished_signal = pyqtSignal(bool)

    def __init__(self, qrcode_key, login_mgr):
        super().__init__()
        self.qrcode_key = qrcode_key
        self.login_mgr = login_mgr
        self._stop = Event()

    def stop(self):
        self._stop.set()

    def run(self):
        import time
        start = time.time()
        while not self._stop.is_set() and time.time() - start < 60:
            success, msg = self.login_mgr.poll_login_status(self.qrcode_key, timeout=2)
            if success:
                self.login_mgr.save_cookies()
                self.finished_signal.emit(True)
                return
            elif msg == "expired":
                self.status_signal.emit("二维码已过期，请重新生成")
                self.finished_signal.emit(False)
                return
            else:
                self.status_signal.emit(msg if msg else "等待扫码...")
            self.msleep(2000)
        self.finished_signal.emit(False)


# ---------- 下载工作线程 ----------
class DownloadWorker(QThread):
    progress_signal = pyqtSignal(int, int, float)
    page_done_signal = pyqtSignal(int, dict)
    task_done_signal = pyqtSignal(dict)

    def __init__(self, pages, options, parser):
        super().__init__()
        self.pages = pages if isinstance(pages, list) else []
        self.options = options
        self.parser = parser
        self._is_cancelled = False
        self._lock = QMutex()

    def cancel(self):
        with QMutexLocker(self._lock):
            self._is_cancelled = True

    def run(self):
        if not self.pages:
            self.task_done_signal.emit({'results': [], 'cancelled': self._is_cancelled})
            return

        total = len(self.pages)
        results = []
        for idx, page in enumerate(self.pages):
            if self._is_cancelled:
                break

            original_page_num = page.get('page_num', idx + 1)
            page_index_0based = original_page_num - 1

            def progress_cb(pct):
                self.progress_signal.emit(page_index_0based, total, pct)
            self.parser.set_progress_callback(progress_cb)

            try:
                res = self.parser.resolve_page(
                    page['bvid'], page['aid'], page['cid'],
                    self.options['quality'],
                    self.options['platform'],
                    self.options['output_dir'],
                    self.options['title'],
                    page_index_0based,
                    self.options['mode'],
                    self.options['keep_intermediate'],
                    self.options['enable_danmaku'],
                    self.options['enable_subtitle'],
                    self.options['subtitle_lans']
                )
                results.append({
                    'page': original_page_num,
                    'status': 'success',
                    'files': res
                })
            except Exception as e:
                results.append({
                    'page': original_page_num,
                    'status': 'error',
                    'error': str(e)
                })
            self.page_done_signal.emit(page_index_0based, results[-1])

        self.task_done_signal.emit({
            'results': results,
            'cancelled': self._is_cancelled
        })


# ---------- 自定义任务项 Widget（使用 QTextEdit 显示详细结果） ----------
class TaskWidget(QWidget):
    def __init__(self, task_id, total_pages, output_dir, parent=None):
        super().__init__(parent)
        self.task_id = task_id
        self.total = total_pages
        self.output_dir = output_dir
        self.page_progress = [0.0] * total_pages if total_pages > 0 else []
        self.results = []

        layout = QVBoxLayout(self)
        layout.setContentsMargins(5, 5, 5, 5)

        top = QHBoxLayout()
        self.id_label = QLabel(f"任务 #{task_id}")
        self.status_label = QLabel("下载中...")
        self.status_label.setStyleSheet("color: blue;")
        top.addWidget(self.id_label)
        top.addStretch()
        top.addWidget(self.status_label)
        layout.addLayout(top)

        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        layout.addWidget(self.progress_bar)

        # ---------- 改用 QTextEdit 显示长文件名 ----------
        self.detail_label = QTextEdit()
        self.detail_label.setReadOnly(True)
        self.detail_label.setMaximumHeight(150)
        self.detail_label.setStyleSheet("background: transparent; border: none;")
        layout.addWidget(self.detail_label)

        self.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.customContextMenuRequested.connect(self.show_context_menu)

    def show_context_menu(self, pos):
        menu = QMenu(self)
        open_action = QAction("打开输出文件夹", self)
        open_action.triggered.connect(self.open_folder)
        menu.addAction(open_action)
        menu.exec(self.mapToGlobal(pos))

    def open_folder(self):
        if os.path.exists(self.output_dir):
            QDesktopServices.openUrl(QUrl.fromLocalFile(self.output_dir))
        else:
            QMessageBox.warning(self, "提示", f"输出目录不存在：{self.output_dir}")

    def update_page_progress(self, page_idx_0based, percent):
        display_page = page_idx_0based + 1
        if 0 <= page_idx_0based < len(self.page_progress):
            self.page_progress[page_idx_0based] = percent
        total_pct = sum(self.page_progress) / self.total if self.total else 0
        self.progress_bar.setValue(int(total_pct))
        self.detail_label.setPlainText(f"分P {display_page} 进度: {percent:.1f}%")

    def add_result(self, result):
        if isinstance(result, dict):
            self.results.append(result)
        else:
            self.results.append({
                'page': len(self.results) + 1,
                'status': 'error',
                'error': f'收到无效结果：{result}'
            })

    def finish(self, summary):
        if not isinstance(self.results, list):
            self.results = []

        cancelled = summary.get('cancelled', False)
        self.status_label.setText("已取消" if cancelled else "已完成")
        self.status_label.setStyleSheet("color: orange;" if cancelled else "color: green;")
        self.progress_bar.setValue(100)

        if self.results:
            text_lines = []
            for r in self.results:
                if not isinstance(r, dict):
                    text_lines.append(f"无效结果: {r}")
                    continue
                if r.get('status') == 'success':
                    files = r.get('files')
                    line = f"P{r.get('page', '?')}: "
                    parts = []
                    if files and isinstance(files, dict):
                        if 'video' in files and files['video']:
                            parts.append(f"视频: {os.path.basename(files['video'])}")
                        if 'audio' in files and files['audio']:
                            parts.append(f"音频: {os.path.basename(files['audio'])}")
                        if 'danmaku' in files and files['danmaku']:
                            parts.append(f"弹幕: {os.path.basename(files['danmaku'])}")
                        if 'subtitles' in files and files['subtitles']:
                            for sub in files['subtitles']:
                                parts.append(f"字幕: {os.path.basename(sub)}")
                    line += " | ".join(parts) if parts else "无文件"
                else:
                    line = f"P{r.get('page', '?')}: ❌ {r.get('error', '未知错误')}"
                text_lines.append(line)
            self.detail_label.setPlainText("\n".join(text_lines))
        else:
            self.detail_label.setPlainText("没有结果数据")


# ---------- 主窗口 ----------
class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.login_mgr = BiliLogin()
        self.parser = None
        self.current_pages = []
        self.current_subtitles = []
        self.tasks = {}
        self.task_counter = 0
        self.poll_thread = None

        self.init_ui()
        self.check_login_status()

    def init_ui(self):
        self.setWindowTitle("B站下载器 GUI (PyQt6)")
        self.setMinimumSize(950, 750)

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

        # ===== 登录区域 =====
        login_group = QGroupBox("登录")
        login_layout = QHBoxLayout()
        self.login_label = QLabel("未登录")
        self.login_label.setStyleSheet("color: gray;")
        btn_cookie = QPushButton("Cookie登录")
        btn_qr = QPushButton("扫码登录")
        btn_logout = QPushButton("退出登录")
        btn_logout.setEnabled(False)
        login_layout.addWidget(self.login_label)
        login_layout.addStretch()
        login_layout.addWidget(btn_cookie)
        login_layout.addWidget(btn_qr)
        login_layout.addWidget(btn_logout)
        login_group.setLayout(login_layout)
        main_layout.addWidget(login_group)

        # ===== 解析区域 =====
        parse_group = QGroupBox("视频解析")
        parse_layout = QVBoxLayout()

        input_row = QHBoxLayout()
        self.url_edit = QLineEdit()
        self.url_edit.setPlaceholderText("请输入BV号、番剧或收藏夹链接...")
        self.parse_btn = QPushButton("解析")
        self.parse_btn.setEnabled(False)
        input_row.addWidget(self.url_edit)
        input_row.addWidget(self.parse_btn)
        parse_layout.addLayout(input_row)

        info_row = QHBoxLayout()
        self.title_label = QLabel("标题：")
        self.owner_label = QLabel("UP主：")
        info_row.addWidget(self.title_label)
        info_row.addWidget(self.owner_label)
        info_row.addStretch()
        parse_layout.addLayout(info_row)

        self.page_table = QTableWidget()
        self.page_table.setColumnCount(3)
        self.page_table.setHorizontalHeaderLabels(["选择", "序号", "标题"])
        self.page_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.Stretch)
        self.page_table.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows)
        parse_layout.addWidget(self.page_table)

        opt_row = QHBoxLayout()
        self.quality_combo = QComboBox()
        for qn, name in sorted(BiliParser.QUALITY_MAP.items(), reverse=True):
            self.quality_combo.addItem(f"{name} ({qn})", qn)
        self.mode_combo = QComboBox()
        self.mode_combo.addItems(["merge", "video", "audio", "separate"])
        self.mode_combo.setCurrentText("merge")
        self.danmaku_check = QCheckBox("弹幕")
        self.subtitle_check = QCheckBox("字幕")
        self.subtitle_lang_combo = QComboBox()
        self.subtitle_lang_combo.addItem("全部", "all")
        self.subtitle_lang_combo.setEnabled(False)

        opt_row.addWidget(QLabel("清晰度："))
        opt_row.addWidget(self.quality_combo)
        opt_row.addWidget(QLabel("模式："))
        opt_row.addWidget(self.mode_combo)
        opt_row.addWidget(self.danmaku_check)
        opt_row.addWidget(self.subtitle_check)
        opt_row.addWidget(QLabel("字幕语言："))
        opt_row.addWidget(self.subtitle_lang_combo)
        opt_row.addStretch()
        parse_layout.addLayout(opt_row)

        dir_row = QHBoxLayout()
        self.dir_edit = QLineEdit("downloads")
        self.dir_btn = QPushButton("浏览...")
        dir_row.addWidget(QLabel("输出目录："))
        dir_row.addWidget(self.dir_edit)
        dir_row.addWidget(self.dir_btn)
        parse_layout.addLayout(dir_row)

        self.download_btn = QPushButton("开始下载")
        self.download_btn.setEnabled(False)
        parse_layout.addWidget(self.download_btn)

        parse_group.setLayout(parse_layout)
        main_layout.addWidget(parse_group)

        # ===== 任务列表 =====
        task_group = QGroupBox("下载任务")
        task_layout = QVBoxLayout()
        self.task_list = QListWidget()
        self.task_list.setMinimumHeight(200)
        task_layout.addWidget(self.task_list)
        task_group.setLayout(task_layout)
        main_layout.addWidget(task_group)

        # ===== 信号连接 =====
        btn_cookie.clicked.connect(self.on_cookie_login)
        btn_qr.clicked.connect(self.on_qr_login)
        btn_logout.clicked.connect(self.on_logout)
        self.parse_btn.clicked.connect(self.on_parse)
        self.download_btn.clicked.connect(self.on_download)
        self.dir_btn.clicked.connect(self.on_browse_dir)
        self.subtitle_check.stateChanged.connect(self.on_subtitle_check)
        self.url_edit.returnPressed.connect(self.on_parse)

    # ---------- 登录相关 ----------
    def check_login_status(self):
        if self.login_mgr.load_cookies():
            self.update_login_status(True)
        else:
            self.update_login_status(False)

    def update_login_status(self, logged_in):
        if logged_in:
            try:
                resp = self.login_mgr.session.get('https://api.bilibili.com/x/web-interface/nav')
                data = resp.json()
                uname = data.get('data', {}).get('uname', '未知用户')
                self.login_label.setText(f"已登录：{uname}")
                self.login_label.setStyleSheet("color: green;")
                self.parse_btn.setEnabled(True)
                self.parser = BiliParser(self.login_mgr.session,
                                         self.login_mgr.img_key,
                                         self.login_mgr.sub_key)
            except:
                self.login_label.setText("已登录 (获取用户信息失败)")
                self.login_label.setStyleSheet("color: green;")
                self.parse_btn.setEnabled(True)
        else:
            self.login_label.setText("未登录")
            self.login_label.setStyleSheet("color: gray;")
            self.parse_btn.setEnabled(False)
            self.download_btn.setEnabled(False)

    def on_cookie_login(self):
        text, ok = QInputDialog.getMultiLineText(self, "Cookie登录",
            "请从浏览器复制完整的 Cookie 字符串（多行）：", "")
        if ok and text.strip():
            cookie_str = text.strip().replace('\n', ';')
            cookies = {}
            for item in cookie_str.split(';'):
                if '=' in item:
                    k, v = item.split('=', 1)
                    cookies[k.strip()] = v.strip()
            if not cookies:
                QMessageBox.warning(self, "错误", "无效的Cookie格式")
                return
            self.login_mgr.session.cookies.update(cookies)
            if self.login_mgr.check_login():
                self.login_mgr.cookies = cookies
                self.login_mgr.save_cookies()
                self.update_login_status(True)
            else:
                QMessageBox.warning(self, "失败", "Cookie无效或已过期")

    def on_qr_login(self):
        try:
            img_bytes, qrcode_key = self.login_mgr.generate_qrcode_bytes()
            pixmap = QPixmap()
            pixmap.loadFromData(img_bytes)

            dialog = QDialog(self)
            dialog.setWindowTitle("扫码登录")
            layout = QVBoxLayout(dialog)
            label_img = QLabel()
            label_img.setPixmap(pixmap.scaled(200, 200, Qt.AspectRatioMode.KeepAspectRatio))
            layout.addWidget(label_img, alignment=Qt.AlignmentFlag.AlignCenter)

            self.qr_status_label = QLabel("请使用B站App扫描二维码")
            layout.addWidget(self.qr_status_label, alignment=Qt.AlignmentFlag.AlignCenter)

            btn_cancel = QPushButton("取消")
            btn_cancel.clicked.connect(dialog.reject)
            layout.addWidget(btn_cancel)

            dialog.setModal(False)
            dialog.show()

            self.poll_thread = PollThread(qrcode_key, self.login_mgr)
            self.poll_thread.status_signal.connect(self.qr_status_label.setText)
            self.poll_thread.finished_signal.connect(
                lambda success: self.on_poll_finished(success, dialog)
            )
            self.poll_thread.start()

        except Exception as e:
            QMessageBox.critical(self, "错误", f"生成二维码失败：{e}")

    def on_poll_finished(self, success, dialog):
        dialog.accept()
        if success:
            self.update_login_status(True)
            QMessageBox.information(self, "成功", "登录成功！")
        else:
            QMessageBox.warning(self, "失败", "登录超时或取消")

    def on_logout(self):
        if os.path.exists("bilibili_cookies.json"):
            os.remove("bilibili_cookies.json")
        self.login_mgr.session.cookies.clear()
        self.update_login_status(False)
        QMessageBox.information(self, "已退出", "已清除登录信息")

    # ---------- 解析 ----------
    def on_parse(self):
        url = self.url_edit.text().strip()
        if not url:
            return
        if not self.parser:
            QMessageBox.warning(self, "提示", "请先登录")
            return

        try:
            id_type, value = parse_url_input(url)
            pages = []
            title = ""
            owner = ""

            if id_type == 'bv':
                info = self.parser.get_video_info(value)
                title = info['title']
                owner = info['owner']['name']
                aid = info['aid']
                for idx, p in enumerate(info['pages'], start=1):
                    pages.append({
                        'bvid': value,
                        'aid': aid,
                        'cid': p['cid'],
                        'part': p['part'],
                        'page': p['page'],
                        'page_num': idx
                    })
            elif id_type in ('ep', 'ss'):
                season_info = get_season_info(self.login_mgr.session,
                                              epid=value if id_type=='ep' else None,
                                              ssid=value if id_type=='ss' else None)
                title = season_info['title']
                owner = season_info.get('actors', '未知')
                episodes = []
                if 'episodes' in season_info:
                    episodes = season_info['episodes']
                elif 'section' in season_info and season_info['section']:
                    for sec in season_info['section']:
                        episodes.extend(sec.get('episodes', []))
                page_counter = 1
                for ep in episodes:
                    ep_aid = ep.get('aid', 0)
                    if 'pages' in ep and ep['pages']:
                        for page in ep['pages']:
                            pages.append({
                                'bvid': ep['bvid'],
                                'aid': ep_aid,
                                'cid': page['cid'],
                                'part': ep['title'] + ' - ' + page['part'],
                                'page': page['page'],
                                'page_num': page_counter
                            })
                            page_counter += 1
                    else:
                        pages.append({
                            'bvid': ep['bvid'],
                            'aid': ep_aid,
                            'cid': ep['cid'],
                            'part': ep.get('long_title') or ep.get('title', '未命名'),
                            'page': 1,
                            'page_num': page_counter
                        })
                        page_counter += 1
            elif id_type == 'fav':
                fav_list = get_fav_list(self.login_mgr.session, value)
                if not fav_list:
                    QMessageBox.warning(self, "提示", "收藏夹为空")
                    return
                title = "收藏夹视频"
                owner = fav_list[0]['upper']['name']
                for idx, item in enumerate(fav_list, start=1):
                    pages.append({
                        'bvid': item['bvid'],
                        'aid': item.get('id', 0),
                        'cid': item['ugc']['first_cid'],
                        'part': item['title'],
                        'page': 1,
                        'page_num': idx
                    })
            else:
                QMessageBox.warning(self, "错误", f"不支持的链接类型：{id_type}")
                return

            if not pages:
                QMessageBox.warning(self, "提示", "未找到任何分P")
                return

            self.current_pages = pages
            self.title_label.setText(f"标题：{title}")
            self.owner_label.setText(f"UP主：{owner}")
            self.display_pages(pages)
            self.download_btn.setEnabled(True)

            if pages:
                try:
                    subs = self.parser.get_subtitle_list_for_page(pages[0]['aid'], pages[0]['cid'])
                    self.current_subtitles = subs
                    self.populate_subtitle_langs(subs)
                except:
                    self.current_subtitles = []
                    self.subtitle_lang_combo.clear()
                    self.subtitle_lang_combo.addItem("无字幕", "all")
                    self.subtitle_lang_combo.setEnabled(False)

        except Exception as e:
            QMessageBox.critical(self, "解析失败", str(e))

    def display_pages(self, pages):
        self.page_table.setRowCount(len(pages))
        self.page_table.setColumnWidth(0, 40)
        self.page_table.setColumnWidth(1, 50)
        for i, p in enumerate(pages):
            cb = QCheckBox()
            cb.setChecked(True)
            self.page_table.setCellWidget(i, 0, cb)
            item = QTableWidgetItem(str(i+1))
            item.setTextAlignment(Qt.AlignmentFlag.AlignCenter)
            self.page_table.setItem(i, 1, item)
            self.page_table.setItem(i, 2, QTableWidgetItem(p['part']))

    def populate_subtitle_langs(self, subtitles):
        self.subtitle_lang_combo.clear()
        if not subtitles:
            self.subtitle_lang_combo.addItem("无字幕", "all")
            self.subtitle_lang_combo.setEnabled(False)
            return
        self.subtitle_lang_combo.addItem("全部", "all")
        for sub in subtitles:
            lan = sub.get('lan', 'unknown')
            doc = sub.get('lan_doc', lan)
            self.subtitle_lang_combo.addItem(f"{doc} ({lan})", lan)
        self.subtitle_lang_combo.setEnabled(True)

    def on_subtitle_check(self, state):
        self.subtitle_lang_combo.setEnabled(state == Qt.CheckState.Checked)

    def on_browse_dir(self):
        dir_ = QFileDialog.getExistingDirectory(self, "选择输出目录", self.dir_edit.text())
        if dir_:
            self.dir_edit.setText(dir_)

    # ---------- 下载 ----------
    def on_download(self):
        selected = []
        for row in range(self.page_table.rowCount()):
            cb = self.page_table.cellWidget(row, 0)
            if cb and cb.isChecked():
                selected.append(self.current_pages[row])

        if not selected:
            QMessageBox.warning(self, "提示", "请至少选择一个分P")
            return

        if not isinstance(selected, list):
            selected = []

        output_dir = os.path.abspath(self.dir_edit.text())

        subtitle_lans = None
        if self.subtitle_check.isChecked():
            lang_data = self.subtitle_lang_combo.currentData()
            if lang_data == "all":
                subtitle_lans = "all"
            else:
                subtitle_lans = [lang_data]

        options = {
            'quality': self.quality_combo.currentData(),
            'platform': 'web',
            'mode': self.mode_combo.currentText(),
            'keep_intermediate': False,
            'enable_danmaku': self.danmaku_check.isChecked(),
            'enable_subtitle': self.subtitle_check.isChecked(),
            'subtitle_lans': subtitle_lans,
            'output_dir': output_dir,
            'title': self.title_label.text().replace("标题：", "").strip()
        }

        self.task_counter += 1
        task_id = self.task_counter
        worker = DownloadWorker(selected, options, self.parser)

        task_widget = TaskWidget(task_id, len(selected), output_dir)
        item = QListWidgetItem(self.task_list)
        item.setSizeHint(task_widget.sizeHint())
        self.task_list.addItem(item)
        self.task_list.setItemWidget(item, task_widget)

        worker.progress_signal.connect(task_widget.update_page_progress)
        worker.page_done_signal.connect(lambda idx, result: task_widget.add_result(result))
        worker.task_done_signal.connect(
            lambda summary: self.on_task_done(task_id, summary, task_widget)
        )
        worker.finished.connect(lambda: self.tasks.pop(task_id, None))
        worker.start()
        self.tasks[task_id] = worker

        self.download_btn.setEnabled(False)

    def on_task_done(self, task_id, summary, task_widget):
        task_widget.finish(summary)
        self.download_btn.setEnabled(True)
        if not summary.get('cancelled'):
            QMessageBox.information(self, "完成", f"任务 #{task_id} 下载完成！")


# ---------- 启动 ----------
if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())