#!/usr/bin/env python3

import argparse
import os
from datetime import datetime
from enum import Enum
import time

# 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 = os.path.join(
    os.environ.get("XDG_PICTURES_DIR", os.path.expanduser("~/Pictures")),
    "Screenshots"
)


def wait_until_file_exists(filepath: str, timeout: int = 5):
    """Wait until a file exists or timeout."""
    start_time = time.time()
    while not os.path.isfile(filepath):
        if time.time() - start_time > timeout:
            return False
        time.sleep(0.1)
    return True


def take_screenshot(typeStr: str, filepath: 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}")
        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 + os.sep + ".niri_screenshot.png"
        if os.path.isfile(niriScreenshotPath):
            os.remove(niriScreenshotPath)
        os.system(cmd[type])
        wait_until_file_exists(niriScreenshotPath)
        if os.path.isfile(niriScreenshotPath):
            os.rename(niriScreenshotPath, filepath)
        wait_until_file_exists(filepath)
    else:
        print("Unsupported desktop environment.")
        exit(1)


def edit_screenshot(filepath: str):
    os.system(f"gradia {SCREENSHOT_DIR}{os.sep}{filepath}")


def file_name(dir, 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
    os.makedirs(SCREENSHOT_DIR, exist_ok=True)
    filename = file_name(SCREENSHOT_DIR)
    filepath = os.path.join(SCREENSHOT_DIR, filename)

    # take screenshot
    take_screenshot(args.type, filepath)

    # check if successful
    if not os.path.isfile(filepath):
        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(filename)
        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()