#!/usr/bin/env python3
import gi
gi.require_version("Gtk", "3.0")
from gi.repository import Gtk, GLib

import json
import subprocess
import threading

FREE_LABEL = {
    "nvidia": "Using X.Org / Nouveau (open source)",
    "broadcom": "Using the open source kernel driver",
}


def list_devices():
    try:
        out = subprocess.run(
            ["/usr/local/sbin/hyprvisor", "list-json"],
            check=True, capture_output=True, text=True,
        ).stdout
        return json.loads(out)
    except Exception:
        return []


class DriverRow(Gtk.Box):
    def __init__(self, device):
        super().__init__(orientation=Gtk.Orientation.VERTICAL, spacing=4)
        self.set_margin_top(10)
        self.set_margin_bottom(10)
        self.set_margin_start(12)
        self.set_margin_end(12)

        self.device = device

        title = Gtk.Label(label=device["label"] or device["package"], xalign=0)
        title.get_style_context().add_class("title")
        self.pack_start(title, False, False, 0)

        self.proprietary = Gtk.RadioButton.new_with_label_from_widget(
            None, f"Using {device['package']} (proprietary, tested)"
        )
        self.free = Gtk.RadioButton.new_with_label_from_widget(
            self.proprietary, FREE_LABEL.get(device["id"], "Using the open source driver")
        )
        self.pack_start(self.proprietary, False, False, 0)
        self.pack_start(self.free, False, False, 0)

        if device["installed"]:
            self.proprietary.set_active(True)
        else:
            self.free.set_active(True)

    def wants_proprietary(self):
        return self.proprietary.get_active()

    def changed(self):
        return self.wants_proprietary() != self.device["installed"]


class MainWindow(Gtk.Window):
    def __init__(self):
        super().__init__(title="Additional Drivers")
        self.set_default_size(560, 400)

        header = Gtk.HeaderBar(title="Additional Drivers")
        header.set_show_close_button(True)
        self.set_titlebar(header)

        self.apply_btn = Gtk.Button(label="Apply Changes")
        self.apply_btn.get_style_context().add_class("suggested-action")
        self.apply_btn.connect("clicked", self.on_apply)
        header.pack_end(self.apply_btn)

        root = Gtk.Box(orientation=Gtk.Orientation.VERTICAL)
        self.add(root)

        self.rows = []
        devices = list_devices()

        if not devices:
            empty = Gtk.Label(
                label="No additional drivers are available for your hardware.",
            )
            empty.set_margin_top(40)
            root.pack_start(empty, True, True, 0)
            self.apply_btn.set_sensitive(False)
        else:
            listbox = Gtk.ListBox()
            listbox.set_selection_mode(Gtk.SelectionMode.NONE)
            for i, device in enumerate(devices):
                if i > 0:
                    listbox.add(Gtk.Separator(orientation=Gtk.Orientation.HORIZONTAL))
                row = DriverRow(device)
                self.rows.append(row)
                listbox.add(row)
            root.pack_start(listbox, True, True, 0)

        self.spinner = Gtk.Spinner()
        self.status_label = Gtk.Label(label="")
        status_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=8)
        status_box.set_margin_top(8)
        status_box.set_margin_bottom(8)
        status_box.set_margin_start(12)
        status_box.set_margin_end(12)
        status_box.pack_start(self.spinner, False, False, 0)
        status_box.pack_start(self.status_label, False, False, 0)
        root.pack_start(status_box, False, False, 0)

    def on_apply(self, _btn):
        changed = [r for r in self.rows if r.changed()]
        if not changed:
            return
        self.apply_btn.set_sensitive(False)
        self.spinner.start()
        self.status_label.set_text("Applying changes...")
        threading.Thread(target=self._apply_all, args=(changed,), daemon=True).start()

    def _apply_all(self, rows):
        ok = True
        for row in rows:
            action = "install" if row.wants_proprietary() else "remove"
            result = subprocess.run(
                ["pkexec", "/usr/local/sbin/hyprvisor", action, row.device["package"]],
            )
            if result.returncode != 0:
                ok = False
        GLib.idle_add(self._apply_done, ok)

    def _apply_done(self, ok):
        self.spinner.stop()
        self.status_label.set_text(
            "Done. Restart your session for changes to take effect."
            if ok else "Applying changes failed -- check the system log."
        )
        self.apply_btn.set_sensitive(False)


def main():
    win = MainWindow()
    win.connect("destroy", Gtk.main_quit)
    win.show_all()
    Gtk.main()


if __name__ == "__main__":
    main()
