#!/usr/bin/env python3
"""Expose-style overview of all i3 workspaces on the focused output.

Click a window tile to focus it, click empty workspace space to switch
to it, Escape/click-outside to cancel. Bind to a key with:

    bindsym $mod+Tab exec --no-startup-id i3-overview
"""
import cairo
import gi

gi.require_version("Gtk", "3.0")
gi.require_version("Gdk", "3.0")
from gi.repository import Gdk, GLib, Gtk
import i3ipc
from Xlib import Xatom
from Xlib import display as xlib_display

ICON_SIZE = 20
_NET_WM_ICON = None
_xdisplay = None


def _get_xdisplay():
    global _xdisplay, _NET_WM_ICON
    if _xdisplay is None:
        _xdisplay = xlib_display.Display()
        _NET_WM_ICON = _xdisplay.intern_atom("_NET_WM_ICON")
    return _xdisplay


def fetch_window_icon(window_id, target_size=ICON_SIZE):
    """Reads _NET_WM_ICON off an X window and returns a cairo ImageSurface
    scaled to target_size, or None if the window has no icon set."""
    if not window_id:
        return None
    try:
        d = _get_xdisplay()
        w = d.create_resource_object("window", window_id)
        prop = w.get_full_property(_NET_WM_ICON, Xatom.CARDINAL)
        if prop is None or not prop.value:
            return None
        data = prop.value

        # _NET_WM_ICON is a concatenation of one or more (w, h, pixels...)
        # blocks; pick the one closest to (but not below) our target size.
        best = None
        i = 0
        n = len(data)
        while i + 2 <= n:
            w_, h_ = int(data[i]), int(data[i + 1])
            count = w_ * h_
            if w_ <= 0 or h_ <= 0 or i + 2 + count > n:
                break
            block = (w_, h_, data[i + 2:i + 2 + count])
            if best is None:
                best = block
            else:
                bw = best[0]
                if (bw < target_size and w_ > bw) or (
                    w_ >= target_size and w_ < bw
                ):
                    best = block
            i += 2 + count
        if best is None:
            return None

        iw, ih, pixels = best
        stride = cairo.ImageSurface.format_stride_for_width(cairo.FORMAT_ARGB32, iw)
        buf = bytearray(stride * ih)
        for y in range(ih):
            row_off = y * stride
            src_off = y * iw
            for x in range(iw):
                argb = int(pixels[src_off + x]) & 0xFFFFFFFF
                a = (argb >> 24) & 0xFF
                r = (argb >> 16) & 0xFF
                g = (argb >> 8) & 0xFF
                b = argb & 0xFF
                # premultiply, cairo ARGB32 native (little-endian BGRA bytes)
                r = r * a // 255
                g = g * a // 255
                b = b * a // 255
                o = row_off + x * 4
                buf[o] = b
                buf[o + 1] = g
                buf[o + 2] = r
                buf[o + 3] = a
        surface = cairo.ImageSurface.create_for_data(
            buf, cairo.FORMAT_ARGB32, iw, ih, stride
        )
        if iw == target_size:
            return surface

        scaled = cairo.ImageSurface(cairo.FORMAT_ARGB32, target_size, target_size)
        cr = cairo.Context(scaled)
        cr.scale(target_size / iw, target_size / ih)
        cr.set_source_surface(surface, 0, 0)
        cr.paint()
        return scaled
    except Exception:
        return None


def fuzzy_match(query, text):
    """Case-insensitive subsequence match: every char of query must appear
    in text, in order, not necessarily contiguous (classic fuzzy-finder)."""
    if not query:
        return True
    if not text:
        return False
    text = text.lower()
    pos = 0
    for ch in query.lower():
        pos = text.find(ch, pos)
        if pos == -1:
            return False
        pos += 1
    return True


MARGIN = 24
GUTTER = 16
WS_TITLE_H = 28
WIN_PADDING = 6
BG = (0.09, 0.09, 0.11, 0.92)
WS_BG = (0.15, 0.15, 0.18, 1.0)
WS_BG_FOCUSED = (0.20, 0.28, 0.38, 1.0)
WIN_BG = (0.24, 0.24, 0.28, 1.0)
WIN_BG_HOVER = (0.32, 0.45, 0.62, 1.0)
WIN_BG_FOCUSED = (0.30, 0.55, 0.85, 1.0)
WIN_BG_FLOAT = (0.30, 0.26, 0.20, 0.95)
WIN_BORDER_FLOAT = (0.85, 0.60, 0.30, 0.9)
TEXT = (0.92, 0.92, 0.94, 1.0)
TEXT_DIM = (0.65, 0.65, 0.7, 1.0)


class Overview(Gtk.Window):
    def __init__(self, i3):
        super().__init__(type=Gtk.WindowType.TOPLEVEL)
        self.i3 = i3
        self.set_decorated(False)
        self.set_app_paintable(True)
        self.set_keep_above(True)
        self.connect("draw", self.on_draw)
        self.connect("button-press-event", self.on_click)
        self.connect("motion-notify-event", self.on_motion)
        self.connect("key-press-event", self.on_key)
        self.add_events(
            Gdk.EventMask.BUTTON_PRESS_MASK
            | Gdk.EventMask.POINTER_MOTION_MASK
        )

        screen = self.get_screen()
        visual = screen.get_rgba_visual()
        if visual:
            self.set_visual(visual)

        tree = i3.get_tree()
        focused = tree.find_focused()
        self.output = focused.workspace().ipc_data["output"] if focused else None
        outputs = {o.name: o for o in i3.get_outputs() if o.active}
        out = outputs.get(self.output) or next(iter(outputs.values()))
        rect = out.rect
        # i3 auto-floats fixed-size windows (min==max size hints); set that
        # up before mapping so it never gets tiled into a workspace.
        self.set_resizable(False)
        self.set_size_request(rect.width, rect.height)
        self.set_default_size(rect.width, rect.height)
        self.set_type_hint(Gdk.WindowTypeHint.DIALOG)
        self.set_skip_taskbar_hint(True)
        self.set_skip_pager_hint(True)
        self.move(rect.x, rect.y)

        self.workspaces = [
            w for w in tree.workspaces() if w.ipc_data.get("output") == out.name
        ]

        self.hit_boxes = []  # (x, y, w, h, kind, target)
        self.hover = None
        self.icon_cache = {}
        self.query = ""

        self.show_all()
        GLib.idle_add(self.place_and_focus, rect.x, rect.y)

    def place_and_focus(self, x, y):
        self.move(x, y)
        self.get_window().focus(Gdk.CURRENT_TIME)
        return False

    def on_key(self, _w, event):
        if event.keyval == Gdk.KEY_Escape:
            if self.query:
                self.query = ""
                self.queue_draw()
            else:
                Gtk.main_quit()
            return True

        if event.keyval == Gdk.KEY_BackSpace:
            self.query = self.query[:-1]
            self.queue_draw()
            return True

        if event.keyval in (Gdk.KEY_Return, Gdk.KEY_KP_Enter):
            matches = self.matching_windows()
            if matches:
                self.i3.command(f'[con_id="{matches[0].id}"] focus')
            Gtk.main_quit()
            return True

        ch = Gdk.keyval_to_unicode(event.keyval)
        if ch and ch >= 32:
            self.query += chr(ch)
            self.queue_draw()
            return True

        return True

    def matching_windows(self):
        result = []
        for ws in self.workspaces:
            for con in ws.leaves():
                if fuzzy_match(self.query, con.name) or fuzzy_match(
                    self.query, con.window_class
                ):
                    result.append(con)
        return result

    def layout_columns(self, n):
        cols = 1
        while cols * cols < n:
            cols += 1
        rows = (n + cols - 1) // cols
        return cols, rows

    def on_draw(self, widget, cr):
        w = widget.get_allocated_width()
        h = widget.get_allocated_height()

        cr.set_source_rgba(*BG)
        cr.paint()

        self.hit_boxes = []

        n = len(self.workspaces)
        if n == 0:
            return

        cols, rows = self.layout_columns(n)
        cell_w = (w - MARGIN * 2 - GUTTER * (cols - 1)) / cols
        cell_h = (h - MARGIN * 2 - GUTTER * (rows - 1)) / rows

        for idx, ws in enumerate(self.workspaces):
            col = idx % cols
            row = idx // cols
            x = MARGIN + col * (cell_w + GUTTER)
            y = MARGIN + row * (cell_h + GUTTER)
            self.draw_workspace(cr, ws, x, y, cell_w, cell_h)

        if self.query:
            self.draw_search_hud(cr, w, h)

    def draw_search_hud(self, cr, w, h):
        label = f"/ {self.query}"
        cr.select_font_face("sans-serif")
        cr.set_font_size(18)
        text_w = cr.text_extents(label)[2]
        box_w = text_w + 32
        box_h = 40
        x = (w - box_w) / 2
        y = h - box_h - 14

        cr.set_source_rgba(0.05, 0.05, 0.06, 0.95)
        self.rounded_rect(cr, x, y, box_w, box_h, 8)
        cr.fill()
        cr.set_source_rgba(*WIN_BORDER_FLOAT)
        self.rounded_rect(cr, x + 1, y + 1, box_w - 2, box_h - 2, 7)
        cr.set_line_width(1.5)
        cr.stroke()

        cr.set_source_rgba(*TEXT)
        cr.move_to(x + 16, y + box_h / 2 + 6)
        cr.show_text(label)

    def draw_workspace(self, cr, ws, x, y, w, h):
        focused = ws.ipc_data.get("focused") or ws.ipc_data.get("visible")
        cr.set_source_rgba(*(WS_BG_FOCUSED if focused else WS_BG))
        self.rounded_rect(cr, x, y, w, h, 10)
        cr.fill()

        cr.set_source_rgba(*TEXT)
        cr.select_font_face("sans-serif")
        cr.set_font_size(15)
        cr.move_to(x + 12, y + 20)
        cr.show_text(ws.name)

        self.hit_boxes.append((x, y, w, WS_TITLE_H, "workspace", ws.name))

        all_leaves = [
            c
            for c in ws.leaves()
            if fuzzy_match(self.query, c.name) or fuzzy_match(self.query, c.window_class)
        ]
        tiled = [c for c in all_leaves if c.floating in (None, "auto_off", "user_off")]
        floating = [c for c in all_leaves if c not in tiled]

        inner_x = x + WIN_PADDING
        inner_y = y + WS_TITLE_H + WIN_PADDING
        inner_w = w - WIN_PADDING * 2
        inner_h = h - WS_TITLE_H - WIN_PADDING * 2

        self.hit_boxes.append((x, inner_y, w, inner_h, "workspace", ws.name))

        if inner_w <= 0 or inner_h <= 0:
            return

        if tiled:
            cols, rows = self.layout_columns(len(tiled))
            tw = (inner_w - GUTTER * (cols - 1) / 2) / cols
            th = (inner_h - GUTTER * (rows - 1) / 2) / rows

            for idx, con in enumerate(tiled):
                col = idx % cols
                row = idx // cols
                wx = inner_x + col * (tw + GUTTER / 2)
                wy = inner_y + row * (th + GUTTER / 2)
                self.draw_window(cr, con, wx, wy, tw, th)

        # Floating windows are drawn on top, scaled to their real position
        # relative to the workspace rect, so they read as floating instead
        # of being just another grid tile.
        ws_rect = ws.ipc_data.get("rect") or {}
        ws_w = ws_rect.get("width") or 1
        ws_h = ws_rect.get("height") or 1
        ws_x0 = ws_rect.get("x", 0)
        ws_y0 = ws_rect.get("y", 0)

        for con in floating:
            fr = con.ipc_data.get("rect") or {}
            rel_x = (fr.get("x", ws_x0) - ws_x0) / ws_w
            rel_y = (fr.get("y", ws_y0) - ws_y0) / ws_h
            rel_w = fr.get("width", ws_w) / ws_w
            rel_h = fr.get("height", ws_h) / ws_h

            fx = inner_x + rel_x * inner_w
            fy = inner_y + rel_y * inner_h
            fw = max(rel_w * inner_w, 24)
            fh = max(rel_h * inner_h, 18)
            self.draw_window(cr, con, fx, fy, fw, fh, floating=True)

    def draw_window(self, cr, con, x, y, w, h, floating=False):
        if w <= 0 or h <= 0:
            return
        is_focused = con.focused
        is_hover = self.hover == con.id
        if floating:
            color = WIN_BG_FOCUSED if is_focused else (WIN_BG_HOVER if is_hover else WIN_BG_FLOAT)
        else:
            color = WIN_BG_FOCUSED if is_focused else (WIN_BG_HOVER if is_hover else WIN_BG)

        if floating:
            cr.set_source_rgba(0, 0, 0, 0.35)
            self.rounded_rect(cr, x + 2, y + 3, w, h, 6)
            cr.fill()

        cr.set_source_rgba(*color)
        self.rounded_rect(cr, x, y, w, h, 6)
        cr.fill()

        if floating:
            cr.set_source_rgba(*WIN_BORDER_FLOAT)
            self.rounded_rect(cr, x + 1, y + 1, w - 2, h - 2, 5)
            cr.set_line_width(1.5)
            cr.stroke()

        icon = self.get_icon(con)
        text_x = x + 8
        if icon is not None and h > ICON_SIZE + 12 and w > ICON_SIZE + 16:
            cr.save()
            cr.translate(x + 8, y + 8)
            cr.set_source_surface(icon, 0, 0)
            cr.paint()
            cr.restore()
            text_x = x + 8 + ICON_SIZE + 6

        title = con.name or con.window_class or "?"
        cr.set_source_rgba(*TEXT)
        cr.select_font_face("sans-serif")
        cr.set_font_size(12)
        self.draw_ellipsized(cr, title, x + 8, y + h / 2 + 4, w - 16)

        if con.window_class:
            cr.set_source_rgba(*TEXT_DIM)
            cr.set_font_size(10)
            self.draw_ellipsized(cr, con.window_class, x + 8, y + h / 2 + 18, w - 16)

        self.hit_boxes.append((x, y, w, h, "window", con.id))

    def get_icon(self, con):
        wid = con.window
        if wid not in self.icon_cache:
            self.icon_cache[wid] = fetch_window_icon(wid)
        return self.icon_cache[wid]

    def draw_ellipsized(self, cr, text, x, y, max_w):
        s = text
        while s and cr.text_extents(s)[2] > max_w:
            s = s[:-1]
        if s != text:
            s = s[:-1] + "…"
        cr.move_to(x, y)
        cr.show_text(s)

    def rounded_rect(self, cr, x, y, w, h, r):
        cr.new_sub_path()
        cr.arc(x + w - r, y + r, r, -1.5708, 0)
        cr.arc(x + w - r, y + h - r, r, 0, 1.5708)
        cr.arc(x + r, y + h - r, r, 1.5708, 3.1416)
        cr.arc(x + r, y + r, r, 3.1416, 4.7124)
        cr.close_path()

    def find_hit(self, mx, my):
        # windows drawn last / on top, so search in reverse for finer hits
        for x, y, w, h, kind, target in reversed(self.hit_boxes):
            if x <= mx <= x + w and y <= my <= y + h:
                return kind, target
        return None, None

    def on_motion(self, widget, event):
        kind, target = self.find_hit(event.x, event.y)
        new_hover = target if kind == "window" else None
        if new_hover != self.hover:
            self.hover = new_hover
            widget.queue_draw()

    def on_click(self, widget, event):
        kind, target = self.find_hit(event.x, event.y)
        if kind == "window":
            self.i3.command(f'[con_id="{target}"] focus')
            Gtk.main_quit()
        elif kind == "workspace":
            self.i3.command(f'workspace "{target}"')
            Gtk.main_quit()
        else:
            Gtk.main_quit()


def main():
    i3 = i3ipc.Connection()
    Overview(i3)
    Gtk.main()


if __name__ == "__main__":
    main()