#!/usr/bin/env python3
"""xdg-desktop-portal-dwm

A minimal xdg-desktop-portal backend for plain X11 window managers
(dwm, i3, awesome, xmonad, ...). Implements the backend side of the
Screenshot portal (org.freedesktop.impl.portal.Screenshot) by shelling
out to maim, so that portal-first apps (Flameshot v14+, Flatpaks,
browsers) work without per-app legacy toggles.

Dependencies:
  - python-dbus-fast (Arch: extra/python-dbus-fast; pip: dbus-fast)
  - maim  (screenshot capture; slop pulled in for -s selection)
  - xdotool + imagemagick (optional, only for PickColor)

Environment overrides (mostly for testing):
  XDPDWM_CAPTURE_CMD              command for full capture, output path appended
  XDPDWM_CAPTURE_INTERACTIVE_CMD  command for interactive capture
  XDPDWM_SAVE_DIR                 where PNGs are written
"""

import asyncio
import os
import shlex
import signal
import subprocess
import sys
import time
from pathlib import Path

from dbus_fast import BusType, Variant
from dbus_fast.aio import MessageBus
from dbus_fast.constants import PropertyAccess
from dbus_fast.service import ServiceInterface, dbus_property, method

BUS_NAME = "org.freedesktop.impl.portal.desktop.dwm"
OBJECT_PATH = "/org/freedesktop/portal/desktop"

# Portal response codes
OK = 0
CANCELLED = 1
ERROR = 2

CAPTURE_TIMEOUT = 120  # seconds; interactive selection can sit a while


def log(*args):
    print("[xdp-dwm]", *args, file=sys.stderr, flush=True)


def save_dir() -> Path:
    d = os.environ.get("XDPDWM_SAVE_DIR")
    if d:
        p = Path(d)
    else:
        # Prefer tmpfs: no disk write, private to the session, cleaned on logout.
        runtime = os.environ.get("XDG_RUNTIME_DIR")
        if runtime and Path(runtime).is_dir():
            p = Path(runtime) / "xdg-desktop-portal-dwm"
        else:
            state = os.environ.get("XDG_STATE_HOME",
                                   os.path.expanduser("~/.local/state"))
            p = Path(state) / "xdg-desktop-portal-dwm"
    p.mkdir(parents=True, exist_ok=True)
    return p


# In-process capture: persistent X connection + XShm grab via mss, minimal
# zlib effort. Avoids fork/exec + maim startup (~100-200ms) per shot.
try:
    import mss as _mss
    import mss.tools as _mss_tools
except ImportError:
    _mss = None

_sct = None  # persistent mss instance (holds the X connection)


def grab_inprocess(out_path: str) -> bool:
    """Full virtual-screen grab via mss. Returns False to fall back to maim."""
    global _sct
    if _mss is None or os.environ.get("XDPDWM_CAPTURE_CMD"):
        return False
    try:
        if _sct is None:
            _sct = _mss.mss()
        mon = _sct.monitors[0]  # entire virtual screen, all heads
        shot = _sct.grab(mon)
        level = int(os.environ.get("XDPDWM_PNG_LEVEL", "1"))
        _mss_tools.to_png(shot.rgb, shot.size, level=level, output=out_path)
        return True
    except Exception as e:
        log("in-process grab failed, falling back to maim:", e)
        try:
            if _sct is not None:
                _sct.close()
        except Exception:
            pass
        _sct = None
        return False


def capture_cmd(interactive: bool, out_path: str) -> list:
    if interactive:
        env = os.environ.get("XDPDWM_CAPTURE_INTERACTIVE_CMD")
        base = shlex.split(env) if env else ["maim", "-s", "-u", "-m", "3"]
    else:
        env = os.environ.get("XDPDWM_CAPTURE_CMD")
        base = shlex.split(env) if env else ["maim", "-u", "-m", "3"]
    return base + [out_path]


async def run(cmd: list, timeout: int = CAPTURE_TIMEOUT):
    """Run a command without blocking the bus. Returns (returncode, stderr)."""
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.DEVNULL,
        stderr=asyncio.subprocess.PIPE,
    )
    try:
        _, err = await asyncio.wait_for(proc.communicate(), timeout=timeout)
    except asyncio.TimeoutError:
        proc.kill()
        return (124, b"timed out")
    return (proc.returncode, err or b"")


class Request(ServiceInterface):
    """org.freedesktop.impl.portal.Request

    The frontend hands us an object path ("handle") for every call and
    may call Close() on it to cancel. Our captures are short-lived, so
    Close() just kills the in-flight capture process if there is one.
    """

    def __init__(self):
        super().__init__("org.freedesktop.impl.portal.Request")
        self.proc = None

    @method()
    def Close(self):
        if self.proc and self.proc.returncode is None:
            try:
                self.proc.kill()
            except ProcessLookupError:
                pass


class ScreenshotPortal(ServiceInterface):
    def __init__(self, bus: MessageBus):
        super().__init__("org.freedesktop.impl.portal.Screenshot")
        self.bus = bus

    @dbus_property(access=PropertyAccess.READ)
    def version(self) -> "u":
        return 3

    @method()
    async def Screenshot(
        self, handle: "o", app_id: "s", parent_window: "s", options: "a{sv}"
    ) -> "ua{sv}":
        interactive = bool(
            options.get("interactive").value if "interactive" in options else False
        )
        log(f"Screenshot request from '{app_id or 'unsandboxed app'}'"
            f" interactive={interactive}")

        req = Request()
        exported = False
        try:
            self.bus.export(handle, req)
            exported = True
        except Exception as e:
            log("could not export Request object (continuing):", e)

        try:
            out = save_dir() / time.strftime("screenshot-%Y%m%d-%H%M%S.png")
            if not interactive:
                t0 = time.monotonic()
                ok = await asyncio.to_thread(grab_inprocess, str(out))
                if ok:
                    log(f"captured {out} in-process "
                        f"({(time.monotonic() - t0) * 1000:.0f}ms)")
                    return [OK, {"uri": Variant("s", out.as_uri())}]
            rc, err = await run(capture_cmd(interactive, str(out)))
            if rc != 0 or not out.exists() or out.stat().st_size == 0:
                out.unlink(missing_ok=True)
                if interactive and rc != 0:
                    # slop selection cancelled with Escape exits non-zero
                    log("interactive capture cancelled/failed:",
                        err.decode(errors="replace").strip())
                    return [CANCELLED, {}]
                log("capture failed:", err.decode(errors="replace").strip())
                return [ERROR, {}]

            log("captured", out)
            return [OK, {"uri": Variant("s", out.as_uri())}]
        except FileNotFoundError as e:
            log("capture tool missing:", e)
            return [ERROR, {}]
        finally:
            if exported:
                try:
                    self.bus.unexport(handle, req)
                except Exception:
                    pass

    @method()
    async def PickColor(
        self, handle: "o", app_id: "s", parent_window: "s", options: "a{sv}"
    ) -> "ua{sv}":
        log(f"PickColor request from '{app_id or 'unsandboxed app'}'")
        try:
            # Let the user click a point: slop -f gives us geometry of a
            # click/selection; fall back to instantaneous cursor position.
            x = y = None
            rc, _ = await run(["which", "slop"], timeout=5)
            if rc == 0:
                p = await asyncio.create_subprocess_exec(
                    "slop", "-f", "%x %y", "-t", "0",
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.DEVNULL,
                )
                out_b, _ = await asyncio.wait_for(p.communicate(), timeout=CAPTURE_TIMEOUT)
                if p.returncode == 0:
                    x, y = out_b.decode().split()
                else:
                    return [CANCELLED, {}]
            else:
                p = await asyncio.create_subprocess_exec(
                    "xdotool", "getmouselocation", "--shell",
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.DEVNULL,
                )
                out_b, _ = await asyncio.wait_for(p.communicate(), timeout=10)
                loc = dict(
                    line.split("=", 1)
                    for line in out_b.decode().strip().splitlines()
                    if "=" in line
                )
                x, y = loc["X"], loc["Y"]

            tmp = save_dir() / f".pick-{os.getpid()}.png"
            rc, err = await run(["maim", "-u", "-g", f"1x1+{x}+{y}", str(tmp)])
            if rc != 0:
                log("pixel grab failed:", err.decode(errors="replace").strip())
                return [ERROR, {}]

            rgb = read_pixel(tmp)
            tmp.unlink(missing_ok=True)
            if rgb is None:
                return [ERROR, {}]
            r, g, b = (c / 255.0 for c in rgb)
            return [OK, {"color": Variant("(ddd)", [r, g, b])}]
        except FileNotFoundError as e:
            log("PickColor helper missing:", e)
            return [ERROR, {}]
        except Exception as e:
            log("PickColor error:", e)
            return [ERROR, {}]


def read_pixel(png_path: Path):
    """Decode the single pixel of a 1x1 PNG. Pure stdlib, no Pillow."""
    import struct
    import zlib

    data = png_path.read_bytes()
    if data[:8] != b"\x89PNG\r\n\x1a\n":
        return None
    pos, idat, meta = 8, b"", {}
    while pos < len(data):
        (length,) = struct.unpack(">I", data[pos:pos + 4])
        ctype = data[pos + 4:pos + 8]
        chunk = data[pos + 8:pos + 8 + length]
        if ctype == b"IHDR":
            w, h, depth, color = struct.unpack(">IIBB", chunk[:10])
            meta = {"w": w, "h": h, "depth": depth, "color": color}
        elif ctype == b"IDAT":
            idat += chunk
        elif ctype == b"IEND":
            break
        pos += 12 + length
    if meta.get("w") != 1 or meta.get("h") != 1 or meta.get("depth") != 8:
        return None
    raw = zlib.decompress(idat)
    # 1 filter byte + pixel bytes
    px = raw[1:]
    if meta["color"] == 2:  # RGB
        return tuple(px[:3])
    if meta["color"] == 6:  # RGBA
        return tuple(px[:3])
    if meta["color"] == 0:  # grayscale
        return (px[0],) * 3
    return None


async def main():
    bus = await MessageBus(bus_type=BusType.SESSION).connect()
    portal = ScreenshotPortal(bus)
    bus.export(OBJECT_PATH, portal)
    await bus.request_name(BUS_NAME)
    log(f"up: {BUS_NAME} at {OBJECT_PATH}")

    stop = asyncio.Event()
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, stop.set)
    await stop.wait()
    log("shutting down")
    bus.disconnect()


if __name__ == "__main__":
    asyncio.run(main())
