#! /usr/bin/python3

import json
import os
import codecs
import setproctitle
import sys
from functools import lru_cache
import gi
gi.require_version('Gtk', '3.0')
gi.require_version('Keybinder', '3.0')
from gi.repository import \
    Gtk, Gdk, GLib, GdkPixbuf, Gio, Keybinder  # noqa: E402

__VERSION__ = "1.0.5"


class Main(object):
    def __init__(self):
        self.zoomlevel = 2
        self.app = Gtk.Application.new(
            "org.kryogenix.magnus",
            Gio.ApplicationFlags.HANDLES_COMMAND_LINE)
        self.app.connect("command-line", self.handle_commandline)
        self.app.connect("shutdown", self.handle_shutdown)
        self.resize_timeout = None
        self.window_metrics = None
        self.window_metrics_restored = False
        self.decorations_height = 0
        self.decorations_width = 0
        self.min_width = 300
        self.min_height = 300
        self.last_x = -1
        self.last_y = -1
        self.refresh_interval = 250
        self.started_by_keypress = False
        self.force_refresh = False
        self.reticule = False
        self.position_lock = False
        self.position_x = 0
        self.position_y = 0
        self.drag_x = 0
        self.drag_y = 0
        self.drag_lock = False
        self.drag_enable = False
        self.max_x = 0
        self.max_y = 0

    def set_max_dimensions(self):
        # ripped directly from https://stackoverflow.com/a/65894179/3569534
        w = 0
        h = 0
        screen = self.w.get_screen().get_display()
        for x in range(0, screen.get_n_monitors()):
            w += screen.get_monitor(x).get_geometry().width
            if (h < screen.get_monitor(x).get_geometry().height):
                h = screen.get_monitor(x).get_geometry().height
        print(f"Max x is {w}, max y is {h}")
        self.max_x = w
        self.max_y = h

    def handle_shutdown(self, app):
        if self.started_by_keypress:
            settings = Gio.Settings.new("org.gnome.desktop.a11y.applications")
            val = settings.get_boolean("screen-magnifier-enabled")
            if val:
                settings.set_boolean("screen-magnifier-enabled", False)

    def handle_commandline(self, app, cmdline):
        args = cmdline.get_arguments()
        if hasattr(self, "w"):
            # already started
            if "--about" in args:
                self.show_about_dialog()
            return 0

        if "--help" in args:
            print("Options:")
            print("  --about")
            print("    Show about dialogue")
            print("  --refresh-interval=120")
            print("    Set refresh interval in milliseconds (lower is faster)")
            print("  --force-refresh")
            print("    Refresh continually (according to refresh interval)")
            print("      even if the mouse has not moved")
            print("  --reticule=255,255,255")
            print("    Use a reticule, and set to this RGB color.")
            print("  --position x y")
            print("    Set center of view, disable mouse tracking")
            print("  --drag")
            print("    Enables drag ability of main view, when --position is used")
            return 0

        if "--force-refresh" in args:
            # If this argument is supplied, refresh the view even if the mouse
            # has not moved. Useful if the screen content is video.
            self.force_refresh = True


        if '--position' in args:
            i = args.index('--position')
            self.position_lock = True
            self.position_x = int(args[i+1])
            self.position_y = int(args[i+2])

        if '--drag' in args:
            self.drag_enable = True

        # Override refresh rate on command line
        for arg in args:
            if arg.startswith("--refresh-interval="):
                parts = arg.split("=")
                if len(parts) == 2:
                    try:
                        rival = int(parts[1])
                        print("Refresh interval set to {}ms".format(rival))
                        self.refresh_interval = rival
                    except ValueError:
                        pass
            if arg == "--started-by-keypress":
                self.started_by_keypress = True
                # This is here so that the autostart desktop file can
                # specify it.
                # The idea is that Gnome-ish desktops have an explicit
                # keybinding to run the system magnifier; what this keybinding
                # actually does, via the {desktop}-settings-daemon, is toggle
                # the gsettings key
                # org.gnome.desktop.a11y.applications screen-magnifier-enabled
                # Magnus provides a desktop file to go in /etc/xdg/autostart
                # which contains an AutostartCondition of
                # GSettings org.gnome.desktop.a11y.applications /
                #       screen-magnifier-enabled
                # and then the {desktop}-session daemon takes care of
                # starting the app when that key goes true, and closing the
                # app if that key goes false. However, the user may also
                # explicitly quit Magnus with the close icon or alt-f4
                # or similar. If they do so, then we explicitly set the key
                # back to false, so that the global keybinding to run the
                # magnifier stays in sync.
            if arg.startswith("--reticule"):
                self.reticule = True
                self.reticule_color = [255, 0, 0]
                try:
                    rgb = arg.split("=")[1].split(",")
                    if len(rgb) == 3:
                        self.reticule_color = [int(i) for i in rgb]
                        print("Using reticule color {}".format(self.reticule_color))
                    else:
                        raise Exception
                except:
                    print("Using fallback reticule color")

        # First time startup
        self.start_everything_first_time()
        if "--about" in args:
            self.show_about_dialog()
        return 0

    def start_everything_first_time(self, on_window_map=None):
        GLib.set_application_name("Magnus")

        # the window
        self.w = Gtk.ApplicationWindow.new(self.app)
        self.w.set_size_request(self.min_width, self.min_height)
        self.w.set_title("Magnus")
        self.w.connect("destroy", lambda a: self.app.quit())
        self.w.connect("configure-event", self.read_window_size)
        self.w.connect("configure-event", self.window_configure)
        self.w.connect("size-allocate", self.read_window_decorations_size)
        devman = self.w.get_screen().get_display().get_device_manager()
        self.pointer = devman.get_client_pointer()
        self.set_max_dimensions()

        # the zoom chooser
        zoom = Gtk.ComboBoxText.new()
        self.zoom = zoom
        for i in range(1, 17):
            zoom.append(str(i), "{}×".format(i))
        zoom.set_active(0)
        zoom.connect("changed", self.set_zoom)

        # RGB label if using a reticule
        if self.reticule:
            self.lbl_color = Gtk.Label.new()

        # the box that contains everything
        self.img = Gtk.Image()
        scrolled_window = Gtk.ScrolledWindow()
        scrolled_window.add(self.img)

        # headerbar or no csd
        use_headerbar = True
        if "GTK_CSD" in os.environ:
            gtk_csd = os.environ.get("GTK_CSD")
            if gtk_csd == "0" or gtk_csd == "no" or gtk_csd == "":
                use_headerbar = False
        if use_headerbar:
            # the headerbar
            head = Gtk.HeaderBar()
            head.set_show_close_button(True)
            head.props.title = "Magnus"
            self.w.set_titlebar(head)
            head.pack_end(zoom)
            if self.reticule:
                head.pack_end(self.lbl_color)
            self.w.add(scrolled_window)
        else:
            # use regular assets
            scrolled_window.set_hexpand(True)
            scrolled_window.set_vexpand(True)
            grid = Gtk.Grid(column_homogeneous=False)
            grid.add(zoom)
            if self.reticule:
                grid.add(self.lbl_color)
            grid.attach(scrolled_window,0,1,4,4)
            self.w.add(grid)

        # bind the zoom keyboard shortcuts
        Keybinder.init()
        if Keybinder.supported():
            Keybinder.bind("<Alt><Super>plus", self.zoom_in, zoom)
            Keybinder.bind("<Alt><Super>equal", self.zoom_in, zoom)
            Keybinder.bind("<Alt><Super>minus", self.zoom_out, zoom)

        # bind mouse drag
        self.w.connect("button-press-event", self.drag_on)
        self.w.connect("button-release-event", self.drag_off)

        # and, go
        self.w.show_all()

        self.width = 0
        self.height = 0
        self.window_x = 0
        self.window_y = 0
        GLib.timeout_add(250, self.read_window_size)

        # and, poll
        GLib.timeout_add(self.refresh_interval, self.poll)

        GLib.idle_add(self.load_config)

    def zoom_out(self, keypress, zoom):
        current_index = zoom.get_active()
        if current_index == 0:
            return
        zoom.set_active(current_index - 1)
        self.set_zoom(zoom)

    def zoom_in(self, keypress, zoom):
        current_index = zoom.get_active()
        size = zoom.get_model().iter_n_children(None)
        if current_index == size - 1:
            return
        zoom.set_active(current_index + 1)
        self.set_zoom(zoom)

    def drag_on(self, a,b):
        if not self.drag_enable:
            return
        display = Gdk.Display.get_default()
        (screen, x, y, modifier) = display.get_pointer()
        self.drag_lock = True
        self.drag_x = x
        self.drag_y = y

    def drag_off(self, a,b):
        self.drag_lock = False

    def read_window_decorations_size(self, win, alloc):
        sz = self.w.get_size()
        self.decorations_width = alloc.width - sz.width
        self.decorations_height = alloc.height - sz.height

    def set_zoom(self, zoom):
        self.zoomlevel = int(zoom.get_active_text().replace("×",""))
        self.poll(force_refresh=True)
        self.serialise()

    def read_window_size(self, *args):
        loc = self.w.get_size()
        self.width = loc.width
        self.height = loc.height

    def show_about_dialog(self, *args):
        about_dialog = Gtk.AboutDialog()
        about_dialog.set_artists(["Stuart Langridge"])
        about_dialog.set_authors(["Stuart Langridge"])
        about_dialog.set_version(__VERSION__)
        about_dialog.set_license_type(Gtk.License.MIT_X11)
        about_dialog.set_website("https://www.kryogenix.org/code/magnus")
        about_dialog.run()
        if about_dialog:
            about_dialog.destroy()

    @lru_cache()
    def makesquares(self, overall_width, overall_height, square_size,
                    value_on, value_off):
        on_sq = list(value_on) * square_size
        off_sq = list(value_off) * square_size
        on_row = []
        off_row = []
        while len(on_row) < overall_width * len(value_on):
            on_row += on_sq
            on_row += off_sq
            off_row += off_sq
            off_row += on_sq
        on_row = on_row[:overall_width * len(value_on)]
        off_row = off_row[:overall_width * len(value_on)]

        on_sq_row = on_row * square_size
        off_sq_row = off_row * square_size

        overall = []
        count = 0
        while len(overall) < overall_width * overall_height * len(value_on):
            overall += on_sq_row
            overall += off_sq_row
            count += 2
        overall = overall[:overall_width * overall_height * len(value_on)]
        return overall

    @lru_cache()
    def get_white_pixbuf(self, width, height):
        square_size = 16
        light = (153, 153, 153, 255)
        dark = (102, 102, 102, 255)
        whole = self.makesquares(width, height, square_size, light, dark)
        arr = GLib.Bytes.new(whole)
        return GdkPixbuf.Pixbuf.new_from_bytes(
            arr, GdkPixbuf.Colorspace.RGB, True, 8,
            width, height, width * len(light))

    def poll(self, force_refresh=False):
        if self.position_lock:

            if self.drag_lock:
            # if dragged, move position lock

                display = Gdk.Display.get_default()
                (screen, x, y, modifier) = display.get_pointer()

                # get delta
                dx = (self.drag_x - x) // self.zoomlevel
                dy = (self.drag_y - y) // self.zoomlevel

                # apply delta
                self.position_x += dx
                self.position_y += dy

                # reset drag start to current position
                self.drag_x = x
                self.drag_y = y

            # use coordinates from position lock
            x,y = self.position_x,self.position_y
            x = min(max(x, 0), self.max_x - 1)
            y = min(max(y, 0), self.max_y - 1)

        else:
        # get new coordinates
            display = Gdk.Display.get_default()
            (screen, x, y, modifier) = display.get_pointer()
            if x == self.last_x and y == self.last_y:
                # bail if nothing would be different
                if not force_refresh and not self.force_refresh:
                    return True
            self.last_x = x
            self.last_y = y

        if (x > self.window_x and
                x <= (self.window_x + self.width + self.decorations_width) and
                y > self.window_y and
                y <= (self.window_y + self.height + self.decorations_height)):
            # pointer is over our window, so make it an empty pixbuf
            white = self.get_white_pixbuf(self.width, self.height)
            self.img.set_from_pixbuf(white)
        else:
            root = Gdk.get_default_root_window()
            scaled_width = self.width // self.zoomlevel
            scaled_height = self.height // self.zoomlevel
            scaled_xoff = scaled_width // 2
            scaled_yoff = scaled_height // 2
            screenshot = Gdk.pixbuf_get_from_window(
                root, x - scaled_xoff,
                y - scaled_yoff, scaled_width, scaled_height)
            scaled_pb = screenshot.scale_simple(
                self.width, self.height,
                GdkPixbuf.InterpType.NEAREST)
            if self.reticule:
                scaled_pb,  color = self.add_reticule(scaled_pb, self.zoomlevel)
                self.lbl_color.set_text(str(color))
            self.img.set_from_pixbuf(scaled_pb)
        return True

    def add_reticule(self, image, zoomlevel):
        """
        Add the reticule to the center.
        """
        # FIXME: the pnm gets distorted if the window is not squared
        w, h = image.get_width(), image.get_height()
        s_w = self.width // self.zoomlevel
        s_h = self.height // self.zoomlevel
        raw_pixels = bytearray(image.get_pixels())
        # We need the remainder, which matters for odd zoomlevel
        z = zoomlevel // 2
        z2 = zoomlevel - z
        cw = w // 2
        ch = h // 2
        startx = cw - (z if s_w % 2 else 0) - 1
        endx = cw + z2 + (0 if s_w % 2 else z)
        starty = ch - (z if s_h % 2 else 0) - 1
        endy = ch + z2 + (0 if s_h % 2 else z)
        # This is not tested against a zoom of 1x, but at least we are
        # collecting the color before drawing the reticule
        center_color = self.get_pixel_color(raw_pixels, startx + 1, starty + 1, w)
        # left and right
        for i in [startx, endx]:
            for j in range(starty, endy + 1):
                raw_pixels = self.draw_pixel(raw_pixels, i, j, self.reticule_color, w, h)
        # top and bottom
        for i in range(startx, endx + 1):
            for j in [starty, endy]:
                raw_pixels = self.draw_pixel(raw_pixels, i, j, self.reticule_color, w, h)
        output_image = GdkPixbuf.PixbufLoader.new_with_type('pnm')
        # magic value for pnm
        output_image.write(bytes(f'P6\n\n{w} {h}\n255\n','ascii'))
        output_image.write(raw_pixels)
        output_image.close()
        output_image = output_image.get_pixbuf()
        return output_image, center_color

    def window_configure(self, window, ev):
        if not self.window_metrics_restored:
            return False
        if self.resize_timeout:
            GLib.source_remove(self.resize_timeout)
        self.resize_timeout = GLib.timeout_add_seconds(
            1, self.save_window_metrics_after_timeout,
            {"x": ev.x, "y": ev.y, "w": ev.width, "h": ev.height})
        self.window_x = ev.x
        self.window_y = ev.y

    def save_window_metrics_after_timeout(self, props):
        GLib.source_remove(self.resize_timeout)
        self.resize_timeout = None
        self.save_window_metrics(props)

    def save_window_metrics(self, props):
        scr = self.w.get_screen()
        sw = float(scr.get_width())
        sh = float(scr.get_height())
        # We save window dimensions as fractions of the screen dimensions,
        # to cope with screen resolution changes while we weren't running
        self.window_metrics = {
            "ww": props["w"] / sw,
            "wh": props["h"] / sh,
            "wx": props["x"] / sw,
            "wy": props["y"] / sh
        }
        self.serialise()

    def restore_window_metrics(self, metrics):
        scr = self.w.get_screen()
        sw = float(scr.get_width())
        sh = float(scr.get_height())
        self.w.set_size_request(self.min_width, self.min_height)
        self.w.resize(
            int(sw * metrics["ww"]), int(sh * metrics["wh"]))
        self.w.move(int(sw * metrics["wx"]), int(sh * metrics["wy"]))

    def get_cache_file(self):
        return os.path.join(GLib.get_user_cache_dir(), "magnus.json")

    def serialise(self, *args, **kwargs):
        # yeah, yeah, supposed to use Gio's async file stuff here. But it was
        # writing corrupted files, and I have no idea why; probably the Python
        # var containing the data was going out of scope or something. Anyway,
        # we're only storing a small JSON file, so life's too short to hammer
        # on this; we'll write with Python and take the hit.
        fp = codecs.open(self.get_cache_file(), encoding="utf8", mode="w")
        data = {"zoom": self.zoomlevel}
        if self.window_metrics:
            data["metrics"] = self.window_metrics
        json.dump(data, fp, indent=2)
        fp.close()

    def load_config(self):
        f = Gio.File.new_for_path(self.get_cache_file())
        f.load_contents_async(None, self.finish_loading_history)

    def finish_loading_history(self, f, res):
        try:
            success, contents, _ = f.load_contents_finish(res)
        except GLib.Error as e:
            print(("couldn't restore settings (error: %s)"
                   ", so assuming they're blank") % (e,))
            contents = "{}"

        try:
            data = json.loads(contents)
        except Exception as e:
            print(("Warning: settings file seemed to be invalid json"
                   " (error: %s), so assuming blank") % (e,))
            data = {}
        zl = data.get("zoom")
        if zl:
            idx = 0
            for row in self.zoom.get_model():
                text, lid = list(row)
                if lid == str(zl):
                    self.zoom.set_active(idx)
                    self.zoomlevel = zl
                idx += 1
        metrics = data.get("metrics")
        if metrics:
            self.restore_window_metrics(metrics)
        self.window_metrics_restored = True

    def draw_pixel(self, old_pixels, x, y, rgb, w, h):
        """
        Draw on the bytearray like it is a canvas.
        Adapted from https://stackoverflow.com/questions/28357155/how-to-edit-pixels-of-a-pixbuf-in-python-gdk
        """
        pixels = old_pixels # bytearray
        # make sure pixel data is reasonable
        x = min(x, w)
        y = min(y, h)
        r = min(rgb[0], 255)
        g = min(rgb[1], 255)
        b = min(rgb[2], 255)
        # insert pixel data at right location in bytes array
        i = y*w + x
        pixels[i*3 + 0] = r
        pixels[i*3 + 1] = g
        pixels[i*3 + 2] = b
        return pixels

    def get_pixel_color(self, pixels, x, y, w):
        """
        Return (r, g, b) tuple for the given pixel from bytearray pixels.
        """
        i = y*w + x
        r, g, b = pixels[i*3:i*3+3]
        return (r, g, b)

def main():
    setproctitle.setproctitle('magnus')
    m = Main()
    m.app.run(sys.argv)


if __name__ == "__main__":
    main()
