import time
import numpy as np
from PIL import Image
import mss
import cv2
import os
from skimage.metrics import structural_similarity as ssim


def take_screenshot():
    with mss.mss() as sct:
        monitor = sct.monitors[0]
        screenshot = sct.grab(monitor)
        img = Image.frombytes("RGB", screenshot.size, screenshot.rgb)
        return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)


def compare_images(img1, img2):
    if img1 is None or img2 is None:
        return 100.0

    # Convert images to grayscale
    gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

    # Compute the Structural Similarity Index (SSIM) between the two images
    score, _ = ssim(gray1, gray2, full=True)

    # The SSIM score is a value between -1 and 1, where 1 means identical image
    # We'll convert it to a percentage difference
    diff_percent = (1 - score) * 100

    return diff_percent


def main():
    last_saved_image = None
    screenshot_count = 0
    if not os.path.exists('screenshots'):
        os.makedirs('screenshots')

    while True:
        current_image = take_screenshot()
        diff_percent = compare_images(current_image, last_saved_image)

        if diff_percent > 15.0:  # Lowered threshold for increased sensitivity
            screenshot_count += 1
            filename = f'screenshots/screenshot_{screenshot_count}.png'
            cv2.imwrite(filename, current_image)
            print(f"Saved {filename} (Difference: {diff_percent:.2f}%)")
            last_saved_image = current_image
        else:
            print(f"Screenshot not saved (Difference: {diff_percent:.2f}%)")

        time.sleep(5)


if __name__ == "__main__":
    main()