import os
import fitz

from PIL import Image
from tkinter import filedialog
from tkinter import Tk, Button, Label, messagebox, Toplevel
from tkinter.filedialog import askopenfilename


def add_image_watermark(input_image_path, watermark_image_path, output_image_path):
    water_width = 100
    water_height = 100
    water_right = 50
    water_bottom = 100
    # 打开原始图片和要添加的水印图片
    original_img = Image.open(input_image_path).convert("RGBA")
    watermark_img = Image.open(watermark_image_path).convert("RGBA")

    # 获取原始图片和水印图片的尺寸
    original_width, original_height = original_img.size

    # 将水印图片缩放到与原始图片相同的大小
    watermark_img = watermark_img.resize((water_width, water_height))

    # 合并原始图片和水印图片，并将结果转换为RGB模式
    # watermarked_img = Image.alpha_composite(original_img, watermark_img).convert("RGB")
    original_img.paste(watermark_img, (original_width - (water_width+water_right),
                       original_height-(water_height+water_bottom)), watermark_img)

    # 保存水印后的图片为JPEG格式
    original_img.save(output_image_path)


class CustomMessageBox(Toplevel):
    def __init__(self, master, text, width=400, height=100):
        super().__init__(master)
        self.transient(master)
        self.title("Error")

        self.label = Label(self, text=text)
        self.label.pack(padx=20, pady=20)

        # 设置了宽和高
        self.width = width
        self.height = height
        self.geometry(
            f'{width}x{height}+{master.winfo_x()}+{master.winfo_y()}')

        self.grab_set()  # 使得对话框不会在点击窗口之外的地方时关闭
        self.wait_window(self)  # 等待对话框关闭


if __name__ == "__main__":
    # items = os.listdir(pdfs)
    # for f in items:
    #     name, suffix = os.path.splitext(f)
    #     if suffix == '.pdf':
    #         out = outPath + '/' + name + '_watered' + suffix
    #         input_file = os.path.join(pdfs, f)
    #         # img2pdf(water_path_png)
    #         # add_watermark(input_file,water_path_pdf,out)
    #         add_watermask_fitz(input_file, water_path_pdf, out)

    # 获取当前目录
    current_dir = os.getcwd()
    # 获取脚本自身的路径
    script_path = os.path.dirname(os.path.abspath(__file__))

    water_path_png = os.path.join(current_dir, 'waters/logo.png')
    water_path_pdf = os.path.join(current_dir, 'waters/logo.pdf')

    outPath = os.path.join(current_dir, 'output')
    pdfs = os.path.join(current_dir, 'pdfs')

    def on_closing():
        root.destroy()
        root.quit()
        exit()

    def select_file(label, extensions=['pdf', 'png', 'jpeg', 'jpg']):
        Tk().withdraw()
        file_path = filedialog.askopenfilename()
        if file_path:
            file_extension = file_path.split('.')[-1].lower()
            if file_extension in extensions:
                label.config(text=file_path)
            else:
                CustomMessageBox(root, "无效的文件类型。请选择PDF、PNG、JPEG或JPG文件。")
                select_file()

    def add_watermask():
        file_path = lable_file.cget("text")
        water_path = lable_stamp.cget("text")
        print(file_path, water_path)
        baseaname = os.path.basename(file_path)
        name, suffix = os.path.splitext(baseaname)
        out_path = os.path.join(os.path.dirname(
            file_path), name + '_watered' + suffix)
        pdf_file = fitz.open(file_path)
        page = pdf_file[0]
        page.insert_image(page.bound(), filename=water_path)
        pdf_file.save(out_path)
        lable_water.config(text=out_path)

    root = Tk()
    root.protocol("WM_DELETE_WINDOW", on_closing)
    root.geometry("600x400")

    # 选择文件
    lable_file = Label(root, text="请选择PDF或图片")
    lable_file.grid(row=0, column=1)
    button_file = Button(root, text="选择PDF",
                         command=lambda: select_file(lable_file))
    button_file.grid(row=0, column=0)

    # 选择公章
    lable_stamp = Label(root, text="请选择选择公章")
    lable_stamp.grid(row=1, column=1)
    button_stamp = Button(root, text="选择公章", command=lambda: select_file(lable_stamp,
                                                                         extensions=['png', 'jpeg', 'jpg']))
    button_stamp.grid(row=1, column=0)

    # 添加水印
    button_water = Button(root, text="添加水印", command=add_watermask)
    button_water.grid(row=2, column=0)
    lable_water = Label(root, text="")
    lable_water.grid(row=2, column=1)

    root.mainloop()