#!/usr/bin/env python3

# Description:
#   A screenshot utility script for Hyprland and Niri desktop environments.
#   Takes screenshots (full, area, window) and provides an option to edit them immediately
#
# Requirements:
# - hyprshot (for Hyprland)
# - grim + slurp (for Niri)
# - gradia (for editing)
# - glib bindings for python

import argparse
import subprocess
import time
import fcntl
from os import environ
from datetime import datetime
from enum import Enum
from pathlib import Path
from shutil import copy2

# 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"
    EDIT = "edit"


SCREENSHOT_DIR = Path.home() / "Pictures" / "Screenshots"


def wait_until_file_exists(filepath: Path, timeout: int = 1):
    """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.05)
    return True


def take_screenshot(filepath: Path, typeStr: str):
    lockFD = open("/tmp/screenshot-script.lock", "w")
    try:
        fcntl.flock(lockFD, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError:
        lockFD.close()
        raise RuntimeError("Another screenshot is currently being taken.")

    try:
        type = ScreenshotType(typeStr)
        currentDesktop = environ.get("XDG_CURRENT_DESKTOP", "")
        if "Hyprland" in currentDesktop:
            cmd = {
                # since I only have one monitor
                ScreenshotType.FULL: f"hyprshot -z -m output -m active -o {SCREENSHOT_DIR} -f ",
                ScreenshotType.AREA: f"hyprshot -z -m region -o {SCREENSHOT_DIR} -f ",
                ScreenshotType.WINDOW: f"hyprshot -z -m window -o {SCREENSHOT_DIR} -f ",
            }
            process = subprocess.run(f"{cmd[type]}{filepath.name}", shell=True)
            if process.returncode != 0:
                raise RuntimeError("Failed to take screenshot: hyprshot command failed.")
            if not wait_until_file_exists(filepath):
                raise RuntimeError("Failed to take screenshot: output file not found after hyprshot command.")

        elif "niri" in currentDesktop:
            niriScreenshotPath = SCREENSHOT_DIR / ".niri_screenshot.png"
            cmd = {
                # niri's built-in screenshot commands are asynchronous, which does not wait for the user to select the area.
                # and the selection ui is drawn inside of niri without its state exposed to external programs.
                # so we use grim + slurp for area mode and niri's built-in commands for others.
                ScreenshotType.FULL: "niri msg action screenshot-screen",
                ScreenshotType.AREA: f" grim -g \"$(slurp)\" -t png {niriScreenshotPath} && cat {niriScreenshotPath} | wl-copy",
                ScreenshotType.WINDOW: "niri msg action screenshot-window",
            }
            if niriScreenshotPath.exists():
                niriScreenshotPath.unlink()
            process = subprocess.run(cmd[type], shell=True)
            if process.returncode != 0:
                print(process.returncode)
                raise RuntimeError("Failed to take screenshot: niri screenshot command failed.")

            if wait_until_file_exists(niriScreenshotPath):
                # niriScreenshotPath.rename(filepath)
                copy2(niriScreenshotPath, filepath)
            else:
                raise RuntimeError("Failed to take screenshot: output file not found after niri command.")
            if not wait_until_file_exists(filepath):
                raise RuntimeError("Failed to take screenshot: output file not found after copying.")

        else:
            # print("Unsupported desktop environment.")
            raise RuntimeError("Unsupported desktop environment.")
    finally:
        fcntl.flock(lockFD, fcntl.LOCK_UN)
        lockFD.close()


def edit_screenshot(filepath: Path):
    subprocess.run(f"gradia {filepath}", shell=True)
    # subprocess.run(f"spectacle -l --edit-existing {filepath}", shell=True)


def gen_file_name(prefix="screenshot", ext=".png"):
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    return f"{prefix}_{timestamp}{ext}"


if __name__ == "__main__":
    Notify.init("Screenshot Utility")

    try:
        # raise RuntimeError("Never gonna give you up, never gonna let you down.")
        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.",
        )
        parser.add_argument(
            "path",
            nargs="?",
            default="",
            help="Path of the given screenshot file (for edit type only).",
        )
        args = parser.parse_args()

        filepath: Path = Path()
        if not args.type == ScreenshotType.EDIT.value:
            # file path
            SCREENSHOT_DIR.mkdir(parents=True, exist_ok=True)
            filename = gen_file_name()
            filepath = SCREENSHOT_DIR / filename

            # take screenshot
            take_screenshot(filepath, args.type)
        else:
            if not args.path:
                raise RuntimeError("Path argument is required for edit type.")
            filepath = Path(args.path).expanduser()
            if not filepath.exists():
                raise RuntimeError(f"File does not exist: {filepath}")

        # create loop instance
        loop = GLib.MainLoop()
        editing = False

        # callback on default action (edit)
        def edit_callback(n, action, user_data):
            try:
                global editing
                editing = True
                edit_screenshot(filepath)
            finally:
                n.close()
                loop.quit()

        # callback on close
        def close_callback(n):
            global editing
            if not editing:
                loop.quit()

        n = Notify.Notification.new(
            # Mako doesn't have action buttons displayed with notification cards,
            "Click to edit",
            str(filepath),
        )
        n.add_action(
            # so default action is used, which will be triggered on simply clicking the notification card
            "default",
            # But for my (or more accurate, Noctalia's) quickshell config, buttons will be displayed with label, even for the default action
            "Open in Editor",
            edit_callback,
            None,
        )

        # set timeout for close_callback
        GLib.timeout_add_seconds(10, close_callback, n)

        n.show()
        loop.run()
    except Exception as e:
        n = Notify.Notification.new(
            "Screenshot Error",
            str(e),
        )
        n.show()