main.py

import sys
import os
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QPushButton, QCheckBox, QFileDialog, 
                             QLabel, QMessageBox, QDoubleSpinBox, QProgressDialog,
                             QStatusBar)
from PyQt6.QtCore import Qt, QMimeData, QUrl
from PyQt6.QtGui import QDragEnterEvent, QDropEvent, QKeySequence, QShortcut, QAction

from widgets.video_player import VideoPlayerWidget
from widgets.timeline import TimelineWidget
from widgets.speed_control import SpeedControlWidget
from utils.export_thread import ExportWorker

def format_time(ms):
    seconds = (ms // 1000) % 60
    minutes = (ms // 60000) % 60
    hours = (ms // 3600000)
    
    if hours > 0:
        return f"{hours:02}:{minutes:02}:{seconds:02}.{ms % 1000:03}"
    return f"{minutes:02}:{seconds:02}.{ms % 1000:03}"

def format_time_short(ms):
    seconds = (ms // 1000) % 60
    minutes = (ms // 60000) % 60
    hours = (ms // 3600000)
    
    if hours > 0:
        return f"{hours:02}:{minutes:02}:{seconds:02}"
    return f"{minutes:02}:{seconds:02}"

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Video Frame Clipper")
        self.resize(1000, 750)
        self.setAcceptDrops(True)
        
        self.current_video_path = None
        self.was_playing_before_drag = False
        
        # Undo/Redo Stacks
        self.undo_stack = []
        self.redo_stack = []
        
        self.init_ui()
        self.apply_styles()
        self.setup_shortcuts()
        
    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        main_layout = QVBoxLayout(central_widget)
        main_layout.setContentsMargins(10, 10, 10, 0) # Bottom 0 for footer
        main_layout.setSpacing(10)
        
        # 1. Video Player Area
        self.video_player = VideoPlayerWidget()
        main_layout.addWidget(self.video_player, stretch=4)
        
        # 2. Timeline Area
        self.timeline = TimelineWidget()
        main_layout.addWidget(self.timeline, stretch=1)
        
        # NEW: Timeline Controls (Play button + Time Label)
        timeline_controls = QHBoxLayout()
        timeline_controls.setContentsMargins(0, 0, 0, 0)
        
        self.btn_play = QPushButton("▶") 
        self.btn_play.setFixedWidth(40)
        self.btn_play.setFocusPolicy(Qt.FocusPolicy.NoFocus)
        self.btn_play.clicked.connect(self.video_player.toggle_play_pause)
        
        self.lbl_time = QLabel("00:00 / 00:00")
        self.lbl_time.setStyleSheet("font-family: monospace; font-weight: bold; font-size: 14px; margin-left: 10px;")
        
        timeline_controls.addWidget(self.btn_play)
        timeline_controls.addWidget(self.lbl_time)
        timeline_controls.addStretch()
        
        main_layout.addLayout(timeline_controls)
        
        # 3. Controls Area
        controls_layout = QHBoxLayout()
        
        self.btn_open = QPushButton("打开视频")
        self.btn_open.setFocusPolicy(Qt.FocusPolicy.NoFocus)
        self.btn_open.clicked.connect(self.open_video_dialog)
        
        speed_label = QLabel("倍速:")
        self.btn_speed = SpeedControlWidget()
        self.btn_speed.speed_changed.connect(self.change_speed)
        
        self.chk_audio = QCheckBox("开启音频")
        self.chk_audio.setFocusPolicy(Qt.FocusPolicy.NoFocus)
        self.chk_audio.setChecked(True) 
        self.chk_audio.toggled.connect(self.toggle_audio)
        
        self.btn_toggle_roi = QPushButton("启用框选")
        self.btn_toggle_roi.setCheckable(True)
        self.btn_toggle_roi.setFocusPolicy(Qt.FocusPolicy.NoFocus)
        self.btn_toggle_roi.toggled.connect(self.toggle_roi_mode)
        
        self.btn_clear_markers = QPushButton("清除标记")
        self.btn_clear_markers.setFocusPolicy(Qt.FocusPolicy.NoFocus)
        self.btn_clear_markers.clicked.connect(self.clear_markers_with_undo)
        
        self.btn_clear_roi = QPushButton("清除选区")
        self.btn_clear_roi.setFocusPolicy(Qt.FocusPolicy.NoFocus)
        self.btn_clear_roi.clicked.connect(self.clear_roi)
        
        self.btn_export = QPushButton("批量导出图片")
        self.btn_export.setFocusPolicy(Qt.FocusPolicy.NoFocus)
        self.btn_export.clicked.connect(self.export_images)
        
        controls_layout.addWidget(self.btn_open)
        controls_layout.addWidget(speed_label)
        controls_layout.addWidget(self.btn_speed)
        controls_layout.addWidget(self.chk_audio)
        controls_layout.addWidget(self.btn_toggle_roi)
        controls_layout.addWidget(self.btn_clear_markers)
        controls_layout.addWidget(self.btn_clear_roi)
        controls_layout.addStretch()
        controls_layout.addWidget(self.btn_export)
        
        main_layout.addLayout(controls_layout)
        
        # 4. Info Footer
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        self.status_bar.showMessage("欢迎使用 | 仓库地址: https://github.com/cniu6/VideoFrameClipper | 使用说明: 空格播放/暂停, F/Enter打点, Ctrl+Z撤销")
        self.status_bar.setStyleSheet("color: #aaaaaa; background-color: #222;")
        
        # Connect signals
        self.video_player.duration_changed.connect(self.timeline.set_duration)
        self.video_player.position_changed.connect(self.timeline.set_position)
        self.video_player.duration_changed.connect(self.update_time_label)
        self.video_player.position_changed.connect(self.update_time_label)
        
        # Play state change to update button icon
        self.video_player.media_player.playbackStateChanged.connect(self.update_play_icon)
        
        # Timeline Dragging Logic
        self.timeline.seek_requested.connect(self.on_seek_start)
        self.timeline.seek_finished.connect(self.on_seek_end)
        
    def setup_shortcuts(self):
        # Undo
        self.shortcut_undo = QShortcut(QKeySequence("Ctrl+Z"), self)
        self.shortcut_undo.activated.connect(self.undo_action)
        
        # Redo (Ctrl+Y or Ctrl+Shift+Z)
        self.shortcut_redo = QShortcut(QKeySequence("Ctrl+Y"), self)
        self.shortcut_redo.activated.connect(self.redo_action)

    def update_time_label(self):
        pos = self.video_player.media_player.position()
        dur = self.video_player.media_player.duration()
        self.lbl_time.setText(f"{format_time_short(pos)} / {format_time_short(dur)}")

    def update_play_icon(self, state):
        if state == self.video_player.media_player.PlaybackState.PlayingState:
            self.btn_play.setText("⏸")
        else:
            self.btn_play.setText("▶")

    def undo_action(self):
        if not self.undo_stack:
            return
            
        action = self.undo_stack.pop()
        command = action['cmd']
        
        if command == 'add_marker':
            # Undo adding: Remove the specific marker
            ts = action['data']
            if ts in self.timeline.markers:
                self.timeline.markers.remove(ts)
                self.redo_stack.append(action)
        
        elif command == 'clear_markers':
            # Undo clear: Restore markers
            old_markers = action['data']
            self.timeline.markers = list(old_markers)
            self.redo_stack.append(action)
            
        self.timeline.update()
        self.show_status(f"撤销操作")

    def redo_action(self):
        if not self.redo_stack:
            return
            
        action = self.redo_stack.pop()
        command = action['cmd']
        
        if command == 'add_marker':
            # Redo adding
            ts = action['data']
            if ts not in self.timeline.markers:
                self.timeline.markers.append(ts)
                self.timeline.markers.sort()
                self.undo_stack.append(action)
                
        elif command == 'clear_markers':
            # Redo clear
            self.undo_stack.append({'cmd': 'clear_markers', 'data': list(self.timeline.markers)})
            self.timeline.markers = []
            
        self.timeline.update()
        self.show_status(f"重做操作")
        
    def add_marker_with_undo(self, timestamp):
        if timestamp not in self.timeline.markers:
            self.timeline.add_marker(timestamp)
            self.undo_stack.append({'cmd': 'add_marker', 'data': timestamp})
            self.redo_stack.clear() # Clear redo on new action
            self.show_status(f"添加标记: {format_time(timestamp)}")

    def clear_markers_with_undo(self):
        if not self.timeline.markers:
            return
        # Save state
        self.undo_stack.append({'cmd': 'clear_markers', 'data': list(self.timeline.markers)})
        self.redo_stack.clear()
        self.clear_markers()
        
    def clear_markers(self):
        self.timeline.markers = []
        self.timeline.update()
        self.show_status("清除所有标记")

    def show_status(self, msg):
        self.status_bar.showMessage(msg, 3000)

    def on_seek_start(self, position):
        if self.video_player.media_player.playbackState() == self.video_player.media_player.PlaybackState.PlayingState:
            self.was_playing_before_drag = True
            self.video_player.media_player.pause()
        self.video_player.media_player.setPosition(position)

    def on_seek_end(self, position):
        self.video_player.media_player.setPosition(position)
        if self.was_playing_before_drag:
            self.video_player.media_player.play()
            self.was_playing_before_drag = False

    def apply_styles(self):
        self.setStyleSheet("""
            QMainWindow { background-color: #2b2b2b; }
            QWidget { color: #ffffff; }
            QPushButton {
                background-color: #3d3d3d;
                border: 1px solid #555;
                padding: 5px 15px;
                border-radius: 3px;
                min-width: 80px;
            }
            QPushButton:hover { background-color: #4d4d4d; }
            QPushButton:pressed { background-color: #2d2d2d; }
            QPushButton:checked { background-color: #555; border: 1px solid #00aaff; }
            QCheckBox { spacing: 5px; }
            QMessageBox { background-color: #2b2b2b; color: #ffffff; }
            QStatusBar { border-top: 1px solid #444; }
        """)

    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            event.accept()
        else:
            event.ignore()

    def dropEvent(self, event: QDropEvent):
        urls = event.mimeData().urls()
        if not urls:
            return
            
        for u in urls:
            file_path = u.toLocalFile()
            if os.path.exists(file_path):
                valid_exts = ('.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv', '.webm', '.m4v')
                if file_path.lower().endswith(valid_exts):
                    self.load_video(file_path)
                    break
                elif file_path.lower().endswith('.lnk'):
                    QMessageBox.information(self, "提示", "暂不支持拖入快捷方式，请拖入视频原文件。")
                    break
            
    def open_video_dialog(self):
        file_name, _ = QFileDialog.getOpenFileName(
            self, 
            "选择视频", 
            "", 
            "Video Files (*.mp4 *.avi *.mov *.mkv *.wmv *.flv *.webm *.m4v)"
        )
        if file_name:
            self.load_video(file_name)
            
    def load_video(self, path):
        self.current_video_path = path
        self.setWindowTitle(f"Video Frame Clipper - {os.path.basename(path)}")
        self.video_player.load_video(path)
        
        # Reset states
        self.timeline.markers = []
        self.timeline.update()
        self.video_player.overlay_widget.clear_selection()
        self.undo_stack.clear()
        self.redo_stack.clear()
        self.btn_play.setText("⏸") # Auto plays
        
    def toggle_audio(self, checked):
        self.video_player.toggle_audio(checked)
        
    def change_speed(self, val):
        self.video_player.set_playback_rate(val)

    def toggle_roi_mode(self, checked):
        self.video_player.overlay_widget.set_drawing_enabled(checked)
        if checked:
            self.btn_toggle_roi.setText("停止框选")
            self.show_status("框选模式开启：拖动画面进行选区")
        else:
            self.btn_toggle_roi.setText("启用框选")
            self.show_status("框选模式关闭")
        
    def clear_roi(self):
        self.video_player.overlay_widget.clear_selection()

    def export_images(self):
        if not self.current_video_path:
            QMessageBox.warning(self, "警告", "请先加载视频！")
            return
            
        if not self.timeline.markers:
            QMessageBox.warning(self, "警告", "没有设置标记点！请按回车键添加红点。")
            return
            
        self.video_player.media_player.pause()
        
        output_dir = os.path.join(os.getcwd(), "out")
        roi_norm = self.video_player.get_roi_normalized()
        
        w = self.video_player.width()
        h = self.video_player.height()
        widget_ar = w / h if h > 0 else 1.0
        
        self.export_worker = ExportWorker(
            self.current_video_path,
            output_dir,
            self.timeline.markers,
            roi_norm,
            widget_ar
        )
        
        self.progress_dialog = QProgressDialog("正在导出图片...", "取消", 0, 100, self)
        self.progress_dialog.setWindowModality(Qt.WindowModality.WindowModal)
        self.progress_dialog.setMinimumDuration(0)
        self.progress_dialog.setAutoClose(True) 
        self.progress_dialog.setAutoReset(True)
        
        self.export_worker.progress_updated.connect(self.progress_dialog.setValue)
        self.export_worker.finished_signal.connect(self.on_export_finished)
        self.progress_dialog.canceled.connect(self.export_worker.cancel)
        
        self.export_worker.start()

    def on_export_finished(self, success, msg, save_dir):
        self.progress_dialog.close()
        
        if success:
            QMessageBox.information(self, "成功", msg)
            try:
                os.startfile(save_dir)
            except Exception as e:
                print(f"Could not open folder: {e}")
        else:
            if "取消" not in msg:
                QMessageBox.critical(self, "提示", msg)

    def keyPressEvent(self, event):
        if event.key() == Qt.Key.Key_Space:
            self.video_player.toggle_play_pause()
        elif event.key() in [Qt.Key.Key_Return, Qt.Key.Key_Enter, Qt.Key.Key_F]:
            current_pos = self.video_player.media_player.position()
            self.add_marker_with_undo(current_pos)
        else:
            super().keyPressEvent(event)
            
    def mousePressEvent(self, event):
        super().mousePressEvent(event)

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