import os
import datetime
from pynput.mouse import Listener
from PIL import ImageGrab, Image
import tkinter as tk

# Function to capture a screenshot
def capture_screenshot(directory, index, cursor_x, cursor_y):
    filename = os.path.join(directory, f"{index}.png")
    img = ImageGrab.grab(bbox=None, include_layered_windows=False, all_screens=True, xdisplay=None)

    # Overlay the cursor image
    cursor_image_path = os.path.join(os.path.expanduser("~"), "Documents", "cursor.png")

    # Open the cursor image
    cursor = Image.open(cursor_image_path)
    cursor = cursor.convert("RGBA")  # Ensure it has an alpha channel for transparency

    # Paste the cursor image onto the screenshot directly at the clicked position
    img.paste(cursor, (cursor_x, cursor_y), cursor)

    img.save(filename)
    print(f"Screenshot saved as {filename}")
    return index + 1

# Function to handle mouse clicks
def on_click(x, y, button, pressed):
    global screenshot_index, save_directory, capturing
    if pressed and capturing:
        screenshot_index = capture_screenshot(save_directory, screenshot_index, x, y)

# Main function to control the screenshot capturing process
def main():
    global screenshot_index, save_directory, capturing, mouse_listener

    # Create a directory for the screenshots
    now = datetime.datetime.now()
    folder_name = now.strftime("%b %d - %I%M%p")
    save_directory = os.path.join(os.path.expanduser("~"), "Documents", folder_name)
    if not os.path.exists(save_directory):
        os.makedirs(save_directory)

    screenshot_index = 1
    capturing = False  # This variable will control whether the script captures screenshots or not

    # Create a UI to control the listener
    root = tk.Tk()
    root.title("Screenshot App")

    def on_start():
        global capturing
        capturing = True
        print("Screenshot capturing started")

    def on_stop():
        global capturing
        capturing = False
        mouse_listener.stop()
        root.quit()

    start_button = tk.Button(root, text="Start", command=on_start)
    stop_button = tk.Button(root, text="Stop", command=on_stop)

    start_button.pack(pady=10)
    stop_button.pack(pady=10)
    root.mainloop()

if __name__ == "__main__":
    mouse_listener = Listener(on_click=on_click)
    mouse_listener.start()
    main()