# coding:utf-8
import os,traceback
from PIL import Image, ImageSequence

#Python PIL 支持的图像文件格式：
#BMP BUFR (identify only) CUR (read only) DCX (read only) EPS (write-only) FITS (identify only) FLI, FLC (read only) FPX (read only) GBR (read only) GD (read only) GIF GRIB (identify only) 
#HDF5 (identify only) ICO (read only) IM IMT (read only) IPTC/NAA (read only) JPEG MCIDAS (read only) MIC (read only) MPEG (identify only) MSP PALM (write only) PCD (read only) PCX PDF (write only) 
#PIXAR (read only) PNG PPM PSD (read only) SGI (read only) SPIDER TGA (read only) TIFF WAL (read only) WMF (identify only) XBM XPM (read only) XV Thumbnails

IMGSuffixList = ['png','jpg','jpeg','bmp','gif']

# Each image add watermask
def waterMark(imgFile, wmFile):
    try:
        orgImg = Image.open(imgFile)
        watermark = Image.open(wmFile)
        imgSize = orgImg.size
        wmSize = watermark.size
        # resize if watermask bigger than org image
        if imgSize[0] < wmSize[0]:
            watermark.resize(tuple(map(lambda x: int(x * 0.5), watermark.size)))
        wmPosition = (imgSize[0]-wmSize[0],imgSize[1]-wmSize[1]) # default right-bottom
        # print(orgImg.format)
        if imgFile.split('.')[-1].upper() == "GIF":
            ims = [m.copy() for m in ImageSequence.Iterator(orgImg)]
            imsNew = []
            for m in ims:
                layer = Image.new('RGBA', m.size)
                layer.paste(watermark, wmPosition)
                imsNew.append(Image.composite(layer, m, layer))

            imsNew[0].save(imgFile, save_all=True, append_images=imsNew[1:])#,duration=1
        else:
            if orgImg.mode == "P":
                orgImg = orgImg.convert('RGB')
            layer = Image.new('RGBA', orgImg.size)
            layer.paste(watermark, wmPosition)
            targetImg = Image.composite(layer, orgImg, layer)
            targetImg.save(imgFile)
    except Exception as e:
        print(imgFile)
        print(traceback.print_exc())

# Batch work
def imageSearch(targetPath,wmFile):
    try:
        fileList = os.listdir(targetPath)
        for file in fileList:
            curPath = os.path.join(targetPath, file)
            if os.path.isdir(curPath):
                imageSearch(curPath,wmFile)
            else:
                if file.split('.')[-1].lower() in IMGSuffixList:
                    waterMark(curPath,wmFile)
    except Exception as e:
        print(traceback.print_exc())


if __name__ =='__main__':
    print("begin...")
    imageSearch("target-dir","watermask-path")
    print("done!!!")