import cv2
import numpy as np
import os
from PIL import Image

def create_watermark(in_img):
    print(in_img.shape, in_img.dtype)
    h, w, _ = in_img.shape
    # Add alpha channel if missing
    if in_img.shape[2] < 4:
        in_img = np.dstack([in_img, np.ones((h, w), dtype="uint8") * 255])

    watermark = cv2.imread("fake.png", cv2.IMREAD_UNCHANGED)
    watermark = cv2.resize(watermark, (288, 191))
    watermark = cv2.cvtColor(watermark, cv2.COLOR_RGBA2BGRA)  # Convert watermark to BGR color space
    wh, ww, _ = watermark.shape
    
    # Tile the watermark to cover the input image
    tiled_watermark = np.tile(watermark, (h // wh + 1, w // ww + 1, 1))[:h, :w, :]

    f1 = np.asarray([0, 0, 0, 250])   # red color filter
    f2 = np.asarray([255, 255, 255, 255])
    mask = cv2.bitwise_not(cv2.inRange(tiled_watermark, f1, f2))
    mask_inv = cv2.bitwise_not(mask)

    res1 = cv2.bitwise_and(in_img, in_img, mask=mask)
    res2 = cv2.bitwise_and(tiled_watermark, tiled_watermark, mask=mask_inv)
    res = cv2.add(res1, res2)

    alpha = 0.5
    return cv2.addWeighted(res, alpha, in_img, 1 - alpha, 0)

def process_images(input_directory, output_directory):
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)
    
    for file in os.listdir(input_directory):
        if file.endswith(".png") or file.endswith(".jpg") or file.endswith(".jpeg"):
            file_path = os.path.join(input_directory, file)
            im = Image.open(file_path)
            res = create_watermark(np.array(im))
            saved = Image.fromarray(res)
            output_path = os.path.join(output_directory, file)
            saved.save(output_path)

if __name__ == "__main__":
    input_directory = './tar'
    output_directory = './tar'
    process_images(input_directory, output_directory)
    
# if __name__ == "__main__":
#     im = Image.open("./try.png")
#     res = create_watermark(np.array(im))
#     saved = Image.fromarray(res)
#     saved.save('./saved.png')

example watermark