better structure

scripts/screenshot

@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+
+import argparse
+import os
+from datetime import datetime
+from enum import Enum
+import time
+from pathlib import Path
+
+# autopep8: off
+import gi
+gi.require_version("Notify", "0.7")
+from gi.repository import Notify, GLib
+# autopep8: on
+
+
+class ScreenshotType(Enum):
+    FULL = "full"
+    AREA = "area"
+    WINDOW = "window"
+
+
+SCREENSHOT_DIR = Path.home() / "Pictures" / "Screenshots"
+
+
+def wait_until_file_exists(filepath: Path, timeout: int = 5):
+    """Wait until a file exists or timeout."""
+    start_time = time.time()
+    while not filepath.exists():
+        if time.time() - start_time > timeout:
+            return False
+        time.sleep(0.1)
+    return True
+
+
+def take_screenshot(filepath: Path, typeStr: str):
+    type = ScreenshotType(typeStr)
+    currentDesktop = os.environ.get("XDG_CURRENT_DESKTOP", "")
+    if "Hyprland" in currentDesktop:
+        cmd = {
+            ScreenshotType.FULL: f"hyprshot -z -m output -m active -o {SCREENSHOT_DIR} -f ",  # since I only have one monitor
+            ScreenshotType.AREA: f"hyprshot -z -m region -o {SCREENSHOT_DIR} -f ",
+            ScreenshotType.WINDOW: f"hyprshot -z -m window -o {SCREENSHOT_DIR} -f ",
+        }
+        os.system(f"{cmd[type]}{filepath.name}")
+        wait_until_file_exists(filepath)
+    elif "niri" in currentDesktop:
+        cmd = {
+            ScreenshotType.FULL: f"niri msg action screenshot-screen",
+            ScreenshotType.AREA: f"niri msg action screenshot",
+            ScreenshotType.WINDOW: f"niri msg action screenshot-window",
+        }
+        niriScreenshotPath = SCREENSHOT_DIR / ".niri_screenshot.png"
+        if niriScreenshotPath.exists():
+            niriScreenshotPath.unlink()
+        os.system(cmd[type])
+        wait_until_file_exists(niriScreenshotPath)
+        if niriScreenshotPath.exists():
+            niriScreenshotPath.rename(filepath)
+        wait_until_file_exists(filepath)
+    else:
+        print("Unsupported desktop environment.")
+        exit(1)
+
+
+def edit_screenshot(filepath: Path):
+    os.system(f"gradia {filepath}")
+
+
+def file_name(dir: Path, prefix="screenshot", ext=".png"):
+    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    return f"{prefix}_{timestamp}{ext}"
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Take screenshots with hyprshot.")
+    parser.add_argument(
+        "type",
+        choices=[t.value for t in ScreenshotType],
+        help="Type of screenshot to take.",
+    )
+    args = parser.parse_args()
+
+    # file path
+    SCREENSHOT_DIR.mkdir(parents=True, exist_ok=True)
+    filename = file_name(SCREENSHOT_DIR)
+    filepath = SCREENSHOT_DIR / filename
+
+    # take screenshot
+    take_screenshot(filepath, args.type)
+
+    # check if successful
+    if not filepath.exists():
+        print("Failed to take screenshot.")
+        exit(1)
+
+    # create loop instance
+    loop = GLib.MainLoop()
+    editing = False
+
+    # callback on default action (edit)
+    def edit_callback(n, action, user_data):
+        global editing
+        editing = True
+        edit_screenshot(filepath)
+        n.close()
+        loop.quit()
+
+    # callback on close
+    def close_callback(n):
+        global editing
+        if not editing:
+            loop.quit()
+
+    # notification
+    Notify.init("Screenshot Utility")
+    n = Notify.Notification.new(
+        "Screenshot taken",
+        "Click to edit"
+    )
+    n.add_action(
+        "default",
+        "Default",
+        edit_callback,
+        None
+    )
+    n.connect("closed", close_callback)
+    n.show()
+    loop.run()