#!/usr/bin/env python3

"""Project sessions: the windows you open together, opened together.

    panama-project save NAME    record the current layout under a name
    panama-project open NAME    lay it out again on free workspaces
    panama-project list         JSON: every project and what it contains
    panama-project delete NAME  forget one

Recorded rather than written by hand. A project is whatever is on screen when
you save it -- which applications, which workspace each was on, and for a
terminal, which directory it was sitting in. That last part is most of the
value: without it a project opens three terminals in your home folder and you
change directory three times.

Workspaces are recorded as positions rather than numbers -- first, second,
third -- and opening a project claims that many free workspaces. A project you
saved from 2, 3 and 4 does not land on top of whatever is on 2, 3 and 4 now.

Windows are launched through Hyprland's Lua dispatch API, because the plain
`hyprctl dispatch exec` form is parsed as Lua on this config and a rules prefix
like [workspace 4 silent] is a syntax error inside it. The same applies to
moving a window by address, which is how an already-running application is
brought into a project rather than launched a second time -- Slack, Thunderbird
and the browser all refuse to open twice and would simply steal focus.
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import sys
import time
from pathlib import Path

PROJECT_DIR = Path(os.environ.get("XDG_DATA_HOME", Path.home() / ".local/share")) / "panama/projects"

# Panama's terminal, matching hypr/keybinds.lua. Opening one in a directory is
# the whole reason a project is worth saving, so the flag matters as much as the
# command.
TERMINAL = "kitty"
TERMINAL_DIRECTORY_FLAG = "-d"

# Applications that refuse to open twice: launching them again focuses the
# existing window and drags you off the workspace you were on. These are moved
# into place instead. Matched against the window class, lowercased.
SINGLE_INSTANCE = re.compile(
    r"slack|thunderbird|helium|discord|spotify|obsidian|bitwarden|steam|firefox|chrom",
    re.IGNORECASE)

NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9 ._-]{0,63}$")


def emit(result: dict) -> int:
    print(json.dumps(result))
    return 0 if result.get("ok") else 2


def hypr(argv: list[str]) -> str:
    try:
        done = subprocess.run(["hyprctl", *argv], capture_output=True,
                              text=True, timeout=15, check=False)
        return done.stdout
    except (OSError, subprocess.SubprocessError):
        return ""


def dispatch(lua: str) -> bool:
    """Run one Lua dispatch through hyprctl eval.

    eval rather than `hyprctl dispatch`: this configuration is Lua, and the
    dispatch subcommand feeds its arguments to the Lua parser, where both
    `[workspace 4 silent] kitty` and `address:0x...` are syntax errors.
    """
    out = hypr(["eval", lua])
    # eval reports failures on stdout and still exits 0, so the exit status
    # cannot be trusted and the text is what says whether it worked.
    return "error" not in out.lower()


def clients() -> list[dict]:
    try:
        parsed = json.loads(hypr(["-j", "clients"]) or "[]")
        return parsed if isinstance(parsed, list) else []
    except json.JSONDecodeError:
        return []


SHELLS = {"bash", "zsh", "fish", "sh", "nu"}


def working_directory(pid: int) -> str | None:
    """Where a terminal is actually sitting, read from the shell inside it.

    Not the terminal's own cwd, which is where the terminal was launched from
    and has nothing to do with where you have since changed to. Reading that
    appeared to work here only by coincidence: these terminals were started from
    a shell that was already in the project directory, so both answers matched.
    One opened from the launcher and then changed into a project would have
    recorded the home folder.

    /proc is the only place this exists -- the compositor knows a window's class
    and title and nothing whatever about what the process inside it is doing.
    """
    # A file of space-separated pids, not a directory. Listing it as one always
    # raised, which meant this quietly fell through to pgrep every single time
    # and the faster path was never once taken.
    try:
        children = Path(f"/proc/{int(pid)}/task/{int(pid)}/children").read_text().split()
    except (OSError, ValueError, TypeError):
        children = []
    if not children:
        try:
            children = subprocess.run(["pgrep", "-P", str(int(pid))],
                                      capture_output=True, text=True,
                                      timeout=5, check=False).stdout.split()
        except (OSError, subprocess.SubprocessError, ValueError):
            return None

    for child in " ".join(children).split():
        try:
            name = Path(f"/proc/{child}/comm").read_text().strip()
            if name not in SHELLS:
                continue
            target = os.readlink(f"/proc/{child}/cwd")
        except (OSError, ValueError):
            continue
        # Home is not worth recording: it is where a terminal opens anyway, and
        # storing it would make every ordinary terminal look deliberate.
        return None if target == str(Path.home()) else target
    return None


# Where desktop entries live, in the order XDG says to prefer them.
APPLICATION_DIRS = [
    Path.home() / ".local/share/applications",
    Path.home() / ".local/share/flatpak/exports/share/applications",
    Path("/var/lib/flatpak/exports/share/applications"),
    Path("/usr/local/share/applications"),
    Path("/usr/share/applications"),
]


def desktop_id(window_class: str) -> str | None:
    """The desktop entry for a window class, if one carries that name.

    Every application on this desktop names its entry after its class --
    kitty.desktop, org.gnome.Nautilus.desktop, com.slack.Slack.desktop -- which
    is what makes this a lookup rather than a guess. Matched case-insensitively,
    because a class is not required to match the filename's case.
    """
    wanted = f"{window_class}.desktop".lower()
    for directory in APPLICATION_DIRS:
        if not directory.is_dir():
            continue
        for entry in directory.glob("*.desktop"):
            if entry.name.lower() == wanted:
                return entry.stem
    return None


# Placeholders a desktop entry uses for the files it was opened with. Nothing
# is being opened here, so they are removed rather than left to arrive as
# literal text. @@u ... @@ is flatpak's file-forwarding wrapper around the same
# idea, and goes with them.
FIELD_CODES = re.compile(r"@@[uUfF]?|\s%[fFuUdDnNickvm]\b")


def desktop_exec(window_class: str) -> str | None:
    """The command a desktop entry says starts this application.

    Read out of the entry rather than handed to gtk-launch, and that is not a
    style preference. gtk-launch activates an application over D-Bus, so the
    process that appears is not a child of the command Hyprland ran -- and a
    [workspace N silent] rule applies to the child. Launching Nautilus that way
    put it on whatever workspace was in front of you instead of the project's.
    Running the Exec line directly keeps it a child, so the rule lands.

    Not /proc/PID/cmdline either: a flatpak reports its own container's path --
    Slack is /app/extra/slack -- which does not exist outside the sandbox, so
    the command that started it cannot start it again. The entry knows the
    `flatpak run` form.
    """
    entry = desktop_id(window_class)
    if entry is None:
        return None
    for directory in APPLICATION_DIRS:
        candidate = directory / f"{entry}.desktop"
        if not candidate.is_file():
            continue
        try:
            for line in candidate.read_text(errors="replace").splitlines():
                if line.startswith("Exec="):
                    command = FIELD_CODES.sub("", line[5:]).strip()
                    return command or None
        except OSError:
            return None
    return None


def path_for(name: str) -> Path:
    return PROJECT_DIR / f"{name}.json"


def save(name: str) -> dict:
    if not NAME_PATTERN.match(name):
        return {"ok": False, "error": "bad-name"}

    # Only real, numbered workspaces. The scratchpad is a place things are put
    # aside, not part of a layout, and special workspaces cannot be recreated by
    # position anyway.
    windows = [w for w in clients()
               if str(w.get("workspace", {}).get("name", "")).isdigit()
               and not w.get("floating", False)]
    if not windows:
        return {"ok": False, "error": "nothing-to-save"}

    # Workspaces are recorded in the order they appear, then referred to by
    # position. What matters is that these three things were on one workspace
    # and those two on the next -- not which numbers they happened to have.
    order: list[str] = []
    grouped: dict[str, list[dict]] = {}
    for window in sorted(windows, key=lambda w: int(w["workspace"]["name"])):
        workspace = window["workspace"]["name"]
        if workspace not in grouped:
            grouped[workspace] = []
            order.append(workspace)

        window_class = str(window.get("class") or window.get("initialClass") or "")
        record: dict = {"class": window_class}

        if window_class == TERMINAL:
            directory = working_directory(window.get("pid", 0))
            if directory:
                record["directory"] = directory

        command = desktop_exec(window_class)
        if command:
            record["exec"] = command

        grouped[workspace].append(record)

    project = {
        "name": name,
        "savedAt": int(time.time()),
        "workspaces": [{"windows": grouped[w]} for w in order],
    }

    PROJECT_DIR.mkdir(parents=True, exist_ok=True)
    path_for(name).write_text(json.dumps(project, indent=2) + "\n")
    return {"ok": True, "name": name,
            "workspaces": len(project["workspaces"]),
            "windows": sum(len(w["windows"]) for w in project["workspaces"])}


def free_workspaces(count: int) -> list[int]:
    """Workspace numbers with nothing on them, lowest first.

    Fresh ones rather than the recorded numbers, so opening a project never
    lands on top of work already in progress. Ten because that is how far the
    keybinds reach -- a project on a workspace ALT cannot select is a project
    that is hard to get back to.
    """
    occupied = {int(w["workspace"]["name"]) for w in clients()
                if str(w.get("workspace", {}).get("name", "")).isdigit()}
    return [n for n in range(1, 11) if n not in occupied][:count]


def running_window(window_class: str) -> str | None:
    """The address of an already-open window of this class, if there is one."""
    for window in clients():
        current = str(window.get("class") or window.get("initialClass") or "")
        if current.lower() == window_class.lower():
            address = window.get("address")
            if isinstance(address, str) and address.startswith("0x"):
                return address
    return None


def place(window: dict, workspace: int) -> str:
    """Get one window onto one workspace. Returns what was done."""
    window_class = str(window.get("class", ""))

    # An application that refuses to open twice is moved rather than launched.
    # Launching it would focus the window it already has and pull the desktop
    # to wherever that was -- the opposite of laying out a project.
    if SINGLE_INSTANCE.search(window_class):
        address = running_window(window_class)
        if address is not None:
            moved = dispatch(
                f'hl.dispatch(hl.dsp.window.move({{ workspace = {workspace}, '
                f'follow = false, window = "address:{address}" }}))')
            return "moved" if moved else "move-failed"

    if window_class == TERMINAL:
        directory = window.get("directory")
        command = TERMINAL
        if isinstance(directory, str) and directory:
            # Quoted for the shell Hyprland runs this through, and only if the
            # directory still exists -- a terminal that fails to open because a
            # project was saved on a machine where that path was there is worse
            # than one that opens in home.
            if Path(directory).is_dir():
                command = f"{TERMINAL} {TERMINAL_DIRECTORY_FLAG} '{directory}'"
    else:
        command = window.get("exec") or window_class

    known = {w.get("address") for w in clients()}
    launched = dispatch(
        f'hl.dispatch(hl.dsp.exec_cmd("[workspace {workspace} silent] {command}"))')
    if not launched:
        return "launch-failed"

    # The rule is not enough on its own. An application marked DBusActivatable
    # -- Nautilus, and most of GNOME -- hands off to its own service and the
    # process Hyprland launched exits immediately, so there is no child left for
    # `[workspace N silent]` to apply to and the window appears on whatever
    # workspace happened to be in front of you.
    #
    # So the window is found afterwards and moved if it needs it. This is also
    # what makes the result honest: a window that never appeared is reported
    # rather than assumed.
    deadline = time.monotonic() + 8
    while time.monotonic() < deadline:
        for window in clients():
            address = window.get("address")
            if address in known or not isinstance(address, str):
                continue
            if str(window.get("class") or "").lower() != window_class.lower():
                continue
            if str(window.get("workspace", {}).get("name", "")) == str(workspace):
                return "launched"
            moved = dispatch(
                f'hl.dispatch(hl.dsp.window.move({{ workspace = {workspace}, '
                f'follow = false, window = "address:{address}" }}))')
            return "launched" if moved else "launched-misplaced"
        time.sleep(0.2)
    return "no-window"


def open_project(name: str) -> dict:
    path = path_for(name)
    if not path.is_file():
        return {"ok": False, "error": "no-such-project"}
    try:
        project = json.loads(path.read_text())
    except json.JSONDecodeError:
        return {"ok": False, "error": "unreadable"}

    workspaces = project.get("workspaces", [])
    targets = free_workspaces(len(workspaces))
    if len(targets) < len(workspaces):
        # Better to say so than to cram two workspaces of a project onto one and
        # leave somebody wondering why it looks wrong.
        return {"ok": False, "error": "not-enough-free-workspaces",
                "needed": len(workspaces), "available": len(targets)}

    placed = []
    for target, workspace in zip(targets, workspaces):
        for window in workspace.get("windows", []):
            placed.append({"class": window.get("class"), "workspace": target,
                           "result": place(window, target)})
            # Hyprland assigns a launched window to the workspace named in the
            # rule, and launching several at once makes that assignment race.
            # A short pause between them is the difference between a laid-out
            # project and a pile on one workspace.
            time.sleep(0.35)

    return {"ok": True, "name": name, "workspaces": targets, "placed": placed}


def list_projects() -> dict:
    projects = []
    if PROJECT_DIR.is_dir():
        for path in sorted(PROJECT_DIR.glob("*.json")):
            try:
                project = json.loads(path.read_text())
            except json.JSONDecodeError:
                continue
            workspaces = project.get("workspaces", [])
            projects.append({
                "name": project.get("name", path.stem),
                "savedAt": project.get("savedAt", 0),
                "workspaces": len(workspaces),
                "windows": sum(len(w.get("windows", [])) for w in workspaces),
                "applications": sorted({str(win.get("class", ""))
                                        for w in workspaces
                                        for win in w.get("windows", [])
                                        if win.get("class")}),
            })
    return {"ok": True, "projects": projects}


def delete_project(name: str) -> dict:
    path = path_for(name)
    if not NAME_PATTERN.match(name) or not path.is_file():
        return {"ok": False, "error": "no-such-project"}
    path.unlink()
    return {"ok": True, "name": name}


def main(argv: list[str]) -> int:
    command = argv[0] if argv else "list"
    argument = argv[1].strip() if len(argv) > 1 else ""

    if command == "list":
        return emit(list_projects())
    if command in {"save", "open", "delete"} and not argument:
        return emit({"ok": False, "error": "no-name"})
    if command == "save":
        return emit(save(argument))
    if command == "open":
        return emit(open_project(argument))
    if command == "delete":
        return emit(delete_project(argument))
    return emit({"ok": False, "error": "unknown-command"})


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
