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


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 Mean Squared Error between the two images
    mse = np.mean((gray1 - gray2) ** 2)

    # Normalize the difference (assuming max pixel value is 255)
    diff_percent = (mse / (255.0 ** 2)) * 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 > 4.0:
            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()