#!/usr/bin/env python3
"""window-restore — bring a minimized window back from the stash workspace.

Mod+M moves the focused window to the declared "minimized" workspace; this is
Mod+Shift+M: list what is stashed in a wofi picker, then move the pick to the
workspace you are currently on and focus it.

Restore targets the *current* workspace by design — the stash workspace is the
entire state, so nothing tracks where a window came from and nothing goes stale.

The stash must be declared (`workspace "minimized"` in niri/config.kdl). niri
silently ignores moves to an undeclared named workspace, so if the declaration
goes missing this script says so rather than failing quietly.
"""
import json
import subprocess
import sys

STASH = "minimized"


# --- pure logic -------------------------------------------------------------
def stash_workspace(workspaces):
    """The declared stash workspace, or None if it is not declared."""
    return next((w for w in workspaces if w.get("name") == STASH), None)


def current_workspace(workspaces):
    return next((w for w in workspaces if w.get("is_focused")), None)


def stashed_windows(windows, stash_id):
    return [w for w in windows if w.get("workspace_id") == stash_id]


def label(window):
    """Display string for the picker. XWayland windows can lack app_id, so fall
    back rather than rendering a bare dash or an empty row."""
    app = (window.get("app_id") or "").strip()
    title = (window.get("title") or "").strip()
    if app and title:
        return f"{app} — {title}"
    return app or title or f"window {window['id']}"


def resolve(windows, chosen):
    """First window whose label matches. Windows with identical labels are
    indistinguishable in the menu anyway, and restoring the first still drains
    the stash on repeat."""
    return next((w for w in windows if label(w) == chosen), None)


def window_exists(windows, window_id):
    """Whether window_id is still present in a fresh `windows` list. Used after
    the picker returns, since wofi can block long enough for the chosen window
    to have closed out from under us."""
    return any(w.get("id") == window_id for w in windows)


def pick(labels, chooser):
    """chooser(text) -> selected line. Injected so tests need no wofi."""
    return chooser("\n".join(labels))


# --- io ---------------------------------------------------------------------
def notify(message):
    subprocess.run(
        ["notify-send", "-a", "window-restore", message],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )


def niri_json(what):
    """`niri msg -j <what>` parsed, or None if niri/IPC is unavailable."""
    try:
        out = subprocess.run(
            ["niri", "msg", "-j", what], capture_output=True, text=True
        )
    except FileNotFoundError:
        return None
    if out.returncode != 0:
        return None
    try:
        return json.loads(out.stdout)
    except json.JSONDecodeError:
        return None


def niri_action(*args):
    subprocess.run(
        ["niri", "msg", "action", *args],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )


def wofi_chooser(text):
    try:
        out = subprocess.run(
            ["wofi", "--dmenu", "--prompt", "Restore:", "--width", "600",
             "--height", "400", "--insensitive"],
            input=text, capture_output=True, text=True,
        )
    except FileNotFoundError:
        return ""
    return out.stdout.strip()


def main():
    workspaces = niri_json("workspaces")
    if workspaces is None:
        notify("niri IPC unavailable")
        return 1

    stash = stash_workspace(workspaces)
    if stash is None:
        notify(f'no "{STASH}" workspace — declare it in niri/config.kdl')
        return 1

    current = current_workspace(workspaces)
    if current is None:
        notify("no focused workspace")
        return 1

    # Restoring here would move a window to the workspace it is already on.
    # Reachable by clicking the waybar pill; no keybind goes here.
    if current["id"] == stash["id"]:
        notify("Already on the minimized workspace")
        return 0

    windows = niri_json("windows")
    if windows is None:
        notify("niri IPC unavailable")
        return 1

    stashed = stashed_windows(windows, stash["id"])
    if not stashed:
        notify("No minimized windows")
        return 0

    chosen = pick([label(w) for w in stashed], wofi_chooser)
    if not chosen:
        return 0  # cancelled

    window = resolve(stashed, chosen)
    if window is None:
        return 0

    # wofi may have sat open for however long the user took. In that time niri
    # can destroy a dynamic workspace (its last window closed) and renumber
    # every later one down, so the pre-pick `current["idx"]` may now name the
    # wrong workspace; the stashed window may also have closed. Re-read state
    # rather than trust anything captured before pick() blocked.
    workspaces = niri_json("workspaces")
    if workspaces is None:
        notify("niri IPC unavailable")
        return 1

    current = current_workspace(workspaces)
    if current is None:
        notify("no focused workspace")
        return 1

    windows = niri_json("windows")
    if windows is None:
        notify("niri IPC unavailable")
        return 1

    if not window_exists(windows, window["id"]):
        notify("that window closed before it could be restored")
        return 0

    # The window is not focused (it is in the stash), so --focus does not apply;
    # focus it explicitly afterwards.
    niri_action("move-window-to-workspace", "--window-id", str(window["id"]),
                str(current["idx"]))
    niri_action("focus-window", "--id", str(window["id"]))
    return 0


if __name__ == "__main__":
    sys.exit(main())