#!/usr/bin/env python3
import subprocess, re

def xrandr(args):
    return subprocess.check_output(["xrandr"] + args).decode("utf-8")

def zenity_list(text, options, height):
    args = ["zenity", "--list", "--hide-header", "--text", text, "--height", str(height), "--column", ""] + list(options)
    res = subprocess.run(args, stdout = subprocess.PIPE)
    if res.returncode == 0:
        output = res.stdout.decode("utf-8")
        if output[-1] == "\n": output = output[:-1]
        return output
    else:
        return None

def current_setup():
    displays = dict()
    walk_display = None

    for line in xrandr(["-q"]).split("\n"):
        if len(line) == 0: continue
        indent = len(re.match(r" *", line)[0])

        if indent == 0:
            if re.match(r"Screen [0-9]+:", line) is not None:
                continue
            else:
                walk_display = line[:line.index(" ")]
                displays[walk_display] = {
                    "current": None,
                    "preferred": None
                }
        elif indent == 2:
            continue
        elif indent == 3:
            words = re.split(r" {2,}", line.strip())
            if len(words) == 0:
                continue

            m = re.match(r"([0-9]+)x([0-9]+)$", words[0])
            resolution = (int(m[1]), int(m[2]))
            freq_current = None  # *
            freq_preferred = None  # +
            for word in words[1:]:
                m = re.match(r"([0-9.]+)([\* ]?)([\+ ]?)$", word)
                if m[2] == "*": freq_current = float(m[1])
                if m[3] == "+": freq_preferred = float(m[1])

            if freq_current is not None:
                displays[walk_display]["current"] = (resolution, freq_current)
            if freq_preferred is not None:
                displays[walk_display]["preferred"] = (resolution, freq_preferred)

    return displays

def iter_find(l, predicate):
    for x in l:
        if predicate(x): return x
    return None

def is_internal_display(name):
    return name.startswith("eDP")

def is_external_display(name):
    return name.startswith("HDMI")

def main():
    setup = current_setup()

    def props_for_name(name):
        nonlocal setup
        current = False if name is None else setup[name]["current"] is not None
        formatted = "<none>" if name is None else name + (" (enabled)" if current else "")
        return {
            "name": name,
            "formatted": formatted,
            "current": current
        }

    internal = props_for_name(iter_find(setup.keys(), is_internal_display))
    external = props_for_name(iter_find(setup.keys(), is_external_display))
    assert internal["name"] is not None
    assert internal["name"] != external["name"]

    other_displays = [setup[name]["formatted"] for name in setup.keys()
                      if name not in [internal["name"], external["name"]]]

    message = "{} displays connected:\n- internal: {}\n- external: {}".format(
                    len(setup.keys()), internal["formatted"], external["formatted"])
    if len(other_displays) > 0:
        message += "\n- unknown: {}".format(other_displays.join(", "))

    options = {
        "Only internal": ["--output", external["name"], "--off"],
        "External left": ["--output", external["name"], "--auto", "--left-of", internal["name"]],
        "External right": ["--output", external["name"], "--auto", "--right-of", internal["name"]],
        "Mirror": ["--output", external["name"], "--auto", "--same-as", internal["name"]]
    }

    response = zenity_list(message, options.keys(), 225)
    if response is not None:
        flags = options[response]
        xrandr(flags)

if __name__ == "__main__":
    main()