#!/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)
# - - (niri has screenshot functionality built-in)
# - gradia (for editing)
# - glib bindings for python

import argparse
import subprocess
import time
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"


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.05)
    return True


def take_screenshot(filepath: Path, typeStr: str):
    type = ScreenshotType(typeStr)
    currentDesktop = 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 ",
        }
        process = subprocess.run(f"{cmd[type]}{filepath.name}", shell=True)
        if process.returncode != 0:
            raise RuntimeError("Failed to take screenshot.")
        wait_until_file_exists(filepath)
    elif "niri" in currentDesktop:
        cmd = {
            ScreenshotType.FULL: "niri msg action screenshot-screen",
            ScreenshotType.AREA: "niri msg action screenshot",
            ScreenshotType.WINDOW: "niri msg action screenshot-window",
        }
        niriScreenshotPath = SCREENSHOT_DIR / ".niri_screenshot.png"
        if niriScreenshotPath.exists():
            niriScreenshotPath.unlink()
        # if os.system(cmd[type]):
        process = subprocess.run(cmd[type], shell=True)
        if process.returncode != 0:
            print("Failed to take screenshot.")
            exit(1)
        wait_until_file_exists(niriScreenshotPath)
        if niriScreenshotPath.exists():
            # niriScreenshotPath.rename(filepath)
            copy2(niriScreenshotPath, filepath)
        else:
            print("Failed to take screenshot.")
            exit(1)
        wait_until_file_exists(filepath)
    else:
        print("Unsupported desktop environment.")
        exit(1)


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


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):
        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()

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

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

    n.show()
    loop.run()
