from watermark import File, Watermark, Position, apply_watermark
import os
from PIL import Image


def add_watermark(src_img_path, watermark_path):
    try:
        # 获取源图片文件名和扩展名
        src_img_name, src_img_ext = os.path.splitext(os.path.basename(src_img_path))

        new_img_ext = src_img_ext
        # 如果源图片是 webp 格式, 先转换成 png
        if src_img_ext.lower() == ".webp":
            tmp_img_path = f"{src_img_name}.png"
            new_img_ext = '.png'
            with Image.open(src_img_path) as img:
                img.save(tmp_img_path, 'PNG')
        else:
            tmp_img_path = src_img_path

        # 添加水印
        watermark = Watermark(File(watermark_path), pos=Position.bottom_left)
        new_file_name = tmp_img_path + new_img_ext
        apply_watermark(File(tmp_img_path), watermark, output_file=new_file_name, overwrite=False)
        
        # 如果源图片是 webp 格式, 转换回 webp
        if src_img_ext.lower() == ".webp":
            img = Image.open(new_file_name)
            img.save(src_img_path, "webp")
            os.remove(tmp_img_path)
            os.remove(new_file_name)
        else:
            os.remove(tmp_img_path)
    except Exception as e:
        print(e)

    
# 使用示例
add_watermark(r'C:\Users\spike\Desktop\xxxxxx\xxxxx.jpg', 'watermark-200.png')