#! /usr/bin/python3
# -*- coding=utf-8 -*-

import gettext
import os, sys
import apt
import subprocess
from gi.repository import GObject, Gdk, Gtk, Gio
from UbuntuDrivers import detect
from aptdaemon import client
from aptdaemon.errors import NotAuthorizedError, TransactionFailed
from aptdaemon.gtk3widgets import AptErrorDialog, AptProgressDialog
import re
import urllib

gettext.install("mintdrivers", "/usr/share/linuxmint/locale")

class Application():

  def __init__(self):

    self.builder = Gtk.Builder()
    self.builder.add_from_file("/usr/lib/linuxmint/mintDrivers/main.ui")
    self.builder.connect_signals(self)
    for o in self.builder.get_objects():
        if issubclass(type(o), Gtk.Buildable):
            name = Gtk.Buildable.get_name(o)
            setattr(self, name, o)
        else:
            print("can not get name for object '%s'" % o)

    self.window_main.show()

    self.window_main.set_title(_("Driver Manager"))

    self.window_main.connect("delete_event", self.quit_application)

    self.button_driver_revert = Gtk.Button(label=_("Re_vert"), use_underline=True)
    self.button_driver_revert.connect("clicked", self.on_driver_changes_revert)
    self.button_driver_apply = Gtk.Button(label=_("_Apply Changes"), use_underline=True)
    self.button_driver_apply.connect("clicked", self.on_driver_changes_apply)
    self.button_driver_cancel = Gtk.Button(label=_("_Cancel"), use_underline=True)
    self.button_driver_cancel.connect("clicked", self.on_driver_changes_cancel)
    self.button_driver_restart = Gtk.Button(label=_("_Restart..."), use_underline=True)
    self.button_driver_restart.connect("clicked", self.on_driver_restart_clicked)
    self.button_driver_revert.set_sensitive(False)
    self.button_driver_revert.set_visible(True)
    self.button_driver_apply.set_sensitive(False)
    self.button_driver_apply.set_visible(True)
    self.button_driver_cancel.set_visible(False)
    self.button_driver_restart.set_visible(False)
    self.box_driver_action.pack_end(self.button_driver_apply, False, False, 0)
    self.box_driver_action.pack_end(self.button_driver_revert, False, False, 0)
    self.box_driver_action.pack_end(self.button_driver_restart, False, False, 0)
    self.box_driver_action.pack_end(self.button_driver_cancel, False, False, 0)

    self.info_bar.set_no_show_all (True)

    self.progress_bar = Gtk.ProgressBar()
    self.box_driver_action.pack_end(self.progress_bar, False, False, 0)
    self.progress_bar.set_visible(False)

    self.needs_restart = False
    self.live_mode = False
    self.apt_client = client.AptClient()

    with open('/proc/cmdline') as f:
      cmdline = f.read()
      if ((not "boot=casper" in cmdline) and (not "boot=live" in cmdline)):
        print ("Post-install mode detected")
        self.info_bar_label.set_text(_("Drivers cannot be installed.\nPlease connect to the Internet or insert the Linux Mint installation DVD (or USB stick)."))
        self.check_internet_or_live_dvd()
        self.button_info_bar.connect("clicked", self.check_internet_or_live_dvd)
      else:
        print ("Live mode detected")
        self.live_mode = True
        self.info_bar.hide()
        self.update_cache()

  def update_cache(self):
    self.update_transaction = self.apt_client.update_cache()
    self.update_transaction.connect("finished", self._on_cache_update_finished)
    dia = AptProgressDialog(self.update_transaction)
    dia.set_transient_for(self.window_main)
    dia.set_modal(True)
    dia.run(close_on_finished=True, show_error=True,
            reply_handler=lambda: True,
            error_handler=self.on_error,
            )

  def on_error(self, error):
    if isinstance(error, NotAuthorizedError):
        # Silently ignore auth failures
        return
    elif not isinstance(error, TransactionFailed):
        # Catch internal errors of the client
        error = TransactionFailed(ERROR_UNKNOWN,str(error))
    dia = AptErrorDialog(error)
    dia.run()
    dia.hide()
    self.apt_cache = apt.Cache()

  def _on_cache_update_finished(self, transaction, exit_state):
    self.show_drivers()

  def quit_application(self, widget=None, event=None):
    self.clean_up_media_cdrom()
    Gtk.main_quit()

  def clean_up_media_cdrom(self):
    if os.path.exists("/media/cdrom"):
      print ("unmounting /media/cdrom...")
      os.system("umount /media/cdrom")
      print ("removing /media/cdrom...")
      os.system("rm -rf /media/cdrom")
      print ("removing cdrom repository...")
      os.system("sed -i '/deb cdrom/d' /etc/apt/sources.list")

  def check_connectivity(self, reference):
    try:
      urllib.request.urlopen(reference, timeout=10)
      return True
    except:
      return False

  def check_internet_or_live_dvd(self, widget=None):

    try:
      print ("Checking connectivity or live media...")

      # We either get drivers from the Internet of from the liveDVD
      # So check the connection to the Internet or scan for a liveDVD

      if self.check_connectivity("http://archive.ubuntu.com"):
        # We can reach the repository, everything's fine
        print ("Internet connection detected")
        self.info_bar.hide()
        self.update_cache()
        return

      # We can't reach the repository, we need the installation media
      mount_point = None
      mounted_on_media_cdrom = False

      # Find the live media
      try:
        live_medias = subprocess.check_output("find /media | grep README.diskdefines", shell=True)
        live_medias = str(live_medias, encoding='utf8').split("\n")
        for live_media in live_medias:
          if ("README.diskdefines" in live_media):
            mount_point = live_media.replace("/README.diskdefines", "") # This is where our live DVD is mounted
            # Add it to apt-cdrom
            os.system("sudo apt-cdrom -d \"%s\" -m add" % mount_point)
            if ("/media/cdrom" in mount_point):
              mounted_on_media_cdrom = True
      except:
        pass

      if mount_point is None:
        # Not mounted..
        self.info_bar.show()
        return

      if mounted_on_media_cdrom == False:
        # Mounted, but not in the right place...
        print ("Binding mount %s to /media/cdrom" % mount_point)
        self.clean_up_media_cdrom()
        os.system("mkdir -p /media/cdrom")
        os.system("mount --bind \"%s\" /media/cdrom" % mount_point)

      # It should now be mounted to /media/cdrom, else something went wrong..
      if os.path.exists("/media/cdrom/README.diskdefines"):
        print ("Live media detected")
        self.info_bar.hide()
        self.update_cache()
      else:
        self.info_bar.show()
    except (NotAuthorizedError, TransactionFailed) as e:
      print("An error occured: {}".format(e))
      self.info_bar.show()

  def on_driver_changes_progress(self, transaction, progress):
    #print(progress)
    self.button_driver_revert.set_visible(False)
    self.button_driver_apply.set_visible(False)
    self.button_driver_restart.set_visible(False)
    self.button_driver_cancel.set_visible(True)
    self.progress_bar.set_visible(True)
    self.progress_bar.set_visible(True)

    self.label_driver_action.set_label(_("Applying changes..."))
    self.progress_bar.set_fraction(progress / 100.0)

  def on_driver_changes_finish(self, transaction, exit_state):
    self.needs_restart = True
    self.progress_bar.set_visible(False)
    self.clear_changes()
    self.apt_cache = apt.Cache()
    self.set_driver_action_status()
    self.update_label_and_icons_from_status()
    self.button_driver_revert.set_visible(True)
    self.button_driver_apply.set_visible(True)
    self.button_driver_cancel.set_visible(False)
    self.scrolled_window_drivers.set_sensitive(True)

  def on_driver_changes_error(self, transaction, error_code, error_details):
    self.on_driver_changes_revert()
    self.set_driver_action_status()
    self.update_label_and_icons_from_status()
    self.button_driver_revert.set_visible(True)
    self.button_driver_apply.set_visible(True)
    self.button_driver_cancel.set_visible(False)
    self.scrolled_window_drivers.set_sensitive(True)
    self.on_error(transaction.error)

  def on_driver_changes_cancellable_changed(self, transaction, cancellable):
    self.button_driver_cancel.set_sensitive(cancellable)

  def on_driver_changes_apply(self, button):

    installs = []
    removals = []

    for pkg in self.driver_changes:
      if pkg.is_installed:
        removals.append(pkg.shortname)
        print ("Remove %s" % pkg.shortname)
      else:
        installs.append(pkg.shortname)
        print ("Install %s" % pkg.shortname)
        try:
          for recommend in pkg.candidate.recommends:
            if recommend[0].name in self.apt_cache:
              installs.append(recommend[0].name)
              print ("Install %s" % recommend[0].name)
        except:
          print ("A problem occured, some recommended deps might not get installed")

    try:
      self.transaction = self.apt_client.commit_packages(install=installs, remove=removals,
                                                         reinstall=[], purge=[], upgrade=[], downgrade=[])
      self.transaction.set_allow_unauthenticated(True)
      self.transaction.connect("progress-changed", self.on_driver_changes_progress)
      self.transaction.connect("cancellable-changed", self.on_driver_changes_cancellable_changed)
      self.transaction.connect("finished", self.on_driver_changes_finish)
      self.transaction.connect("error", self.on_driver_changes_error)
      self.transaction.run()
      self.button_driver_revert.set_sensitive(False)
      self.button_driver_apply.set_sensitive(False)
      self.scrolled_window_drivers.set_sensitive(False)
    except (NotAuthorizedError, TransactionFailed) as e:
      print("Warning: install transaction not completed successfully: {}".format(e))


  def on_driver_changes_revert(self, button_revert=None):

    # HACK: set all the "Do not use" first; then go through the list of the
    #       actually selected drivers.
    for button in self.no_drv:
      button.set_active(True)

    for alias in self.orig_selection:
      button = self.orig_selection[alias]
      button.set_active(True)

    self.clear_changes()

    self.button_driver_revert.set_sensitive(False)
    self.button_driver_apply.set_sensitive(False)

  def on_driver_changes_cancel(self, button_cancel):
    self.transaction.cancel()
    self.clear_changes()

  def on_driver_restart_clicked(self, button_restart):
    self.clean_up_media_cdrom()
    subprocess.call(['reboot'])

  def clear_changes(self):
    self.orig_selection = {}
    self.driver_changes = []

  def on_driver_selection_changed(self, button, modalias, pkg_name=None):
    if self.ui_building:
      return

    pkg = None
    try:
      if pkg_name:
        pkg = self.apt_cache[pkg_name]
    except KeyError:
      pass

    if button.get_active():
      if pkg in self.driver_changes:
        self.driver_changes.remove(pkg)

      if (pkg is not None
         and modalias in self.orig_selection
         and button is not self.orig_selection[modalias]):
        self.driver_changes.append(pkg)
    else:
      if pkg in self.driver_changes:
        self.driver_changes.remove(pkg)

      # for revert; to re-activate the original radio buttons.
      if modalias not in self.orig_selection:
        self.orig_selection[modalias] = button

      if (pkg is not None
         and pkg not in self.driver_changes
         and pkg.is_installed):
        self.driver_changes.append(pkg)

    self.button_driver_revert.set_sensitive(bool(self.driver_changes))
    self.button_driver_apply.set_sensitive(bool(self.driver_changes))

  def gather_device_data(self, device):
    '''Get various device data used to build the GUI.

      return a tuple of (overall_status string, icon, drivers dict).
      the drivers dict is using this form:
        {"recommended/alternative": {pkg_name: {
                                                  'selected': True/False
                                                  'description': 'description'
                                                  'builtin': True/False,
                                                  'free': True/False
                                                }
                                     }}
         "manually_installed": {"manual": {'selected': True, 'description': description_string}}
         "no_driver": {"no_driver": {'selected': True/False, 'description': description_string}}

         Please note that either manually_installed and no_driver are set to None if not applicable
         (no_driver isn't present if there are builtins)
    '''

    possible_overall_status = {
      'recommended': (_("This device is using the recommended driver."), "recommended-driver"),
      'alternative': (_("This device is using an alternative driver."), "other-driver"),
      'manually_installed': (_("This device is using a manually-installed driver."), "other-driver"),
      'no_driver': (_("This device is not working."), "disable-device")
    }

    returned_drivers = {'recommended': {}, 'alternative': {}, 'manually_installed': {}, 'no_driver': {}}
    have_builtin = False
    one_selected = False
    try:
      if device['manual_install']:
        returned_drivers['manually_installed'] = {True: {'selected': True,
                                                  'description': _("Continue using a manually installed driver")}}
    except KeyError:
      pass

    for pkg_driver_name in device['drivers']:
      current_driver = device['drivers'][pkg_driver_name]

      # get general status
      driver_status = 'alternative'
      try:
        if current_driver['recommended'] and current_driver['from_distro']:
          driver_status = 'recommended'
      except KeyError:
        pass

      builtin = False
      try:
        if current_driver['builtin']:
          builtin = True
          have_builtin = True
      except KeyError:
        pass

      try:
        pkg = self.apt_cache[pkg_driver_name]
        installed = pkg.is_installed
        description_line1 = "<b>%s</b>" % pkg.shortname
        description_line2 = "<small><span foreground='#3c3c3c'>%s</span></small> %s" % (_("Version"), pkg.candidate.version)
        description_line3 = "<small><span foreground='#3c3c3c'>%s</span></small>" % pkg.candidate.summary
        if driver_status == 'recommended':
          description_line1 = "%s <b><small><span foreground='#58822B'>(%s)</span></small></b>" % (description_line1, _("recommended"))
        if current_driver['free'] and pkg.shortname != "bcmwl-kernel-source":
          description_line1 = "%s <b><small><span foreground='#2b3882'>(%s)</span></small></b>" % (description_line1, _("open-source"))
        if pkg.shortname == "bcmwl-kernel-source":
          # Always recommend STA driver (B43 is always present alongside it now)
          description_line1 = "%s <b><small><span foreground='#58822B'>(%s)</span></small></b>" % (description_line1, _("recommended"))
        if pkg.shortname.startswith("firmware-b43"):
          # B43 requires a connection to the Internet
          description_line1 = "%s <b><small><span foreground='#822b32'>(%s)</span></small></b>" % (description_line1, _("requires a connection to the Internet"))
        description = "%s\n%s\n%s" % (description_line1, description_line2, description_line3)
      except KeyError:
        print("WARNING: a driver ({}) doesn't have any available package associated: {}".format(pkg_driver_name, current_driver))
        continue

      selected = False
      if not builtin and not returned_drivers['manually_installed']:
        selected = installed
        if installed:
          selected = True
          one_selected = True

      returned_drivers[driver_status].setdefault(pkg_driver_name, {'selected': selected,
                                                                   'description': description,
                                                                   'builtin': builtin,
                                                                   'free': current_driver['free']})

    # adjust making the needed addition
    if not have_builtin:
      selected = False
      if not one_selected:
        selected = True
      returned_drivers["no_driver"] = {True: {'selected': selected,
                                              'description': _("Do not use the device")}}
    else:
      # we have a builtin and no selection: builtin is the selected one then
      if not one_selected:
        for section in ('recommended', 'alternative'):
          for pkg_name in returned_drivers[section]:
            if returned_drivers[section][pkg_name]['builtin']:
              returned_drivers[section][pkg_name]['selected'] = True

    # compute overall status
    for section in returned_drivers:
      for keys in returned_drivers[section]:
        if returned_drivers[section][keys]['selected']:
          (overall_status, icon) = possible_overall_status[section]

    return (overall_status, icon, returned_drivers)

  def get_device_icon(self, device):
    vendor = device.get('vendor', _('Unknown'))
    model =  device.get('model', _('Unknown'))
    icon = "generic"
    if "nvidia" in vendor.lower():
      icon = "nvidia"
    elif "radeon" in vendor.lower() or "radeon" in model.lower() or "Advanced Micro Devices" in vendor or "AMD" in vendor or "ATI" in vendor:
      icon = "ati"
    elif "broadcom" in vendor.lower():
      icon = "broadcom"
    elif "virtualbox" in vendor.lower():
      icon = "virtualbox"
    return ("/usr/share/linuxmint/mintDrivers/icons/%s.png" % icon)


  def show_drivers(self):
    self.apt_cache = apt.Cache()
    self.devices = detect.system_device_drivers()
    self.driver_changes = []
    self.orig_selection = {}
    # HACK: the case where the selection is actually "Do not use"; is a little
    #       tricky to implement because you can't check for whether a package is
    #       installed or any such thing. So let's keep a list of all the
    #       "Do not use" radios, set those active first, then iterate through
    #       orig_selection when doing a Reset.
    self.no_drv = []
    self.nonfree_drivers = 0
    self.ui_building = True
    self.dynamic_device_status = {}
    for device in sorted(self.devices.keys()):
      (overall_status, icon, drivers) = self.gather_device_data(self.devices[device])
      brand_icon = Gtk.Image()
      brand_icon.set_valign(Gtk.Align.START)
      brand_icon.set_halign(Gtk.Align.CENTER)
      brand_icon.set_from_file(self.get_device_icon(self.devices[device]))
      driver_status = Gtk.Image()
      driver_status.set_valign(Gtk.Align.START)
      driver_status.set_halign(Gtk.Align.CENTER)
      driver_status.set_from_icon_name(icon, Gtk.IconSize.MENU)
      device_box = Gtk.Box(spacing=6, orientation=Gtk.Orientation.HORIZONTAL)
      device_box.pack_start(brand_icon, False, False, 6)
      device_detail = Gtk.Box(spacing=6, orientation=Gtk.Orientation.VERTICAL)
      device_box.pack_start(device_detail, True, True, 0)
      model_name = self.devices[device].get('model', None)
      vendor_name = self.devices[device].get('vendor', None)
      if vendor_name is None and model_name is None:
        device_name = _("Unknown")
      elif vendor_name is None:
        device_name = model_name
      elif model_name is None:
        device_name = vendor_name
      else:
        device_name = "%s: %s" % (vendor_name, model_name)
      widget = Gtk.Label(device_name)
      widget.set_halign(Gtk.Align.START)
      device_detail.pack_start(widget, True, False, 0)
      widget = Gtk.Label("<small>{}</small>".format(overall_status))
      widget.set_halign(Gtk.Align.START)
      widget.set_use_markup(True)
      device_detail.pack_start(widget, True, False, 0)
      self.dynamic_device_status[device] = (driver_status, widget)

      option_group = None
      # define the order of introspection
      for section in ('recommended', 'alternative', 'manually_installed', 'no_driver'):
        for driver in sorted(drivers[section], key=lambda x: self.sort_string(drivers[section], x)):
          radio_button = Gtk.RadioButton.new(None)
          label = Gtk.Label()
          label.set_markup(drivers[section][driver]['description'])
          radio_button.add(label)
          if option_group:
            radio_button.join_group(option_group)
          else:
            option_group = radio_button
          device_detail.pack_start(radio_button, True, False, 0)
          radio_button.set_active(drivers[section][driver]['selected'])

          if section == 'no_driver':
            self.no_drv.append(radio_button)
          if section in ('manually_install', 'no_driver') or ('builtin' in drivers[section][driver] and drivers[section][driver]['builtin']):
            radio_button.connect("toggled", self.on_driver_selection_changed, device)
          else:
            radio_button.connect("toggled", self.on_driver_selection_changed, device, driver)
          if drivers['manually_installed'] and section != 'manually_installed' and "firmware" not in str(driver):
            radio_button.set_sensitive(False)

      self.box_driver_detail.pack_start(device_box, False, False, 6)

    self.ui_building = False
    self.box_driver_detail.show_all()
    self.set_driver_action_status()

  def sort_string(self, drivers, x):
    value = drivers[x]['description']
    try:
      value = "%s %s" % (not drivers[x]['free'], value)
    except:
      pass #best effort (some driver options don't have a 'free' flag, and that's alright)
    return value

  def update_label_and_icons_from_status(self):
    '''Update the current label and icon, computing the new device status'''

    for device in self.devices:
      (overall_status, icon, drivers) = self.gather_device_data(self.devices[device])
      (driver_status, widget) = self.dynamic_device_status[device]

      driver_status.set_from_icon_name(icon, Gtk.IconSize.MENU)
      widget.set_label("<small>{}</small>".format(overall_status))


  def set_driver_action_status(self):
    # Update the label in case we end up having some kind of proprietary driver in use.
    if (not self.live_mode) and (os.path.exists('/var/run/reboot-required') or self.needs_restart):
      self.label_driver_action.set_label(_("You need to restart the computer to complete the driver changes."))
      self.button_driver_restart.set_visible(True)
      self.window_main.set_urgency_hint(True)
      return

    self.nonfree_drivers = 0
    for device in self.devices:
      for pkg_name in self.devices[device]['drivers']:
        pkg = self.apt_cache[pkg_name]
        if (not self.devices[device]['drivers'][pkg_name]['free'] or pkg_name == "bcmwl-kernel-source") and pkg.is_installed:
          self.nonfree_drivers = self.nonfree_drivers + 1

    if self.nonfree_drivers > 0:
      self.label_driver_action.set_label(gettext.ngettext (
                                                "%(count)d proprietary driver in use.",
                                                "%(count)d proprietary drivers in use.",
                                                self.nonfree_drivers)
                                                % { 'count': self.nonfree_drivers})
    else:
      self.label_driver_action.set_label(_("No proprietary drivers are in use."))

if __name__ == "__main__":
    if os.getuid() != 0:
        os.execvp("gksu", ("", " ".join(sys.argv)))
    else:
        Application()
        Gtk.main()
