import sys
from PyQt5.QtWidgets import (
    QApplication, QWidget, QPushButton, QVBoxLayout, QLabel, QFileDialog, QMessageBox
)
from moviepy.editor import VideoFileClip
import cv2
import numpy as np
import os

def remove_watermark_from_frame(frame, region):
    # region: (x, y, w, h)
    x, y, w, h = region
    roi = frame[y:y+h, x:x+w]
    # 使用inpaint方法去除水印
    mask = np.ones(roi.shape[:2], np.uint8) * 255
    inpainted = cv2.inpaint(roi, mask, 3, cv2.INPAINT_TELEA)
    frame[y:y+h, x:x+w] = inpainted
    return frame

class WatermarkRemover(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.video_path = None

    def initUI(self):
        self.setWindowTitle('视频2水印去除工具')
        self.layout = QVBoxLayout()

        self.label = QLabel('请选择一个视频文件')
        self.layout.addWidget(self.label)

        self.btn_open = QPushButton('打开视频')
        self.btn_open.clicked.connect(self.open_video)
        self.layout.addWidget(self.btn_open)

        self.btn_remove = QPushButton('去除水印')
        self.btn_remove.setEnabled(False)
        self.btn_remove.clicked.connect(self.process_video)
        self.layout.addWidget(self.btn_remove)

        self.setLayout(self.layout)

    def open_video(self):
        path, _ = QFileDialog.getOpenFileName(self, '选择视频文件', '', 'Video Files (*.mp4 *.avi *.mov)')
        if path:
            self.video_path = path
            self.label.setText(f'已选择: {os.path.basename(path)}')
            self.btn_remove.setEnabled(True)

    def process_video(self):
        if not self.video_path:
            return

        # 这里假设水印在右下角，区域为宽高均50像素（可根据实际情况调整）
        region = self.get_watermark_region(self.video_path)
        output_path = self.video_path.rsplit('.', 1)[0] + '_no_watermark.mp4'

        clip = VideoFileClip(self.video_path)
        w, h = clip.size

        def process_frame(frame):
            # frame: RGB, uint8
            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            result_bgr = remove_watermark_from_frame(frame_bgr, region)
            return cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)

        new_clip = clip.fl_image(process_frame)
        new_clip.write_videofile(output_path, audio=True, threads=4, preset='ultrafast')
        QMessageBox.information(self, '完成', f'处理完成，文件已保存为:\n{output_path}')

    def get_watermark_region(self, video_path):
        # 可根据实际水印位置调整
        # 这里只是一个例子：右下角50x50
        clip = VideoFileClip(video_path)
        w, h = clip.size
        region_w, region_h = 50, 50
        x = w - region_w - 10  # 离右边10像素
        y = h - region_h - 10  # 离下边10像素
        return (x, y, region_w, region_h)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    remover = WatermarkRemover()
    remover.show()
    sys.exit(app.exec_())