#!/usr/bin/python
#
#       switchlibglx
#
#       Copyright 2011 Canonical Ltd.
#       Author: Alberto Milone <alberto.milone@canonical.com>
#
#       Script to switch between AMD and Intel graphics driver libraries.
#
#       Usage:
#           switchlibglx   amd|intel|query
#           amd:   switches to AMD's version of libglx.so
#           intel: switches to the open-source version of libglx.so
#           query: checks which version is currently active and writes
#                  "amd", "intel" or "unknown" to the standard output
#
#       Permission is hereby granted, free of charge, to any person
#       obtaining a copy of this software and associated documentation
#       files (the "Software"), to deal in the Software without
#       restriction, including without limitation the rights to use,
#       copy, modify, merge, publish, distribute, sublicense, and/or sell
#       copies of the Software, and to permit persons to whom the
#       Software is furnished to do so, subject to the following
#       conditions:
#
#       The above copyright notice and this permission notice shall be
#       included in all copies or substantial portions of the Software.
#
#       THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#       EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
#       OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#       NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
#       HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
#       WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#       FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
#       OTHER DEALINGS IN THE SOFTWARE.


import optparse
import os
import sys
import re
import subprocess
from subprocess import Popen, PIPE, CalledProcessError

class Alternatives:

    def __init__(self, master_link):
        self.__open_drivers_alternative = 'mesa/ld.so.conf'
        self.__command = 'update-alternatives'
        self.__master_link = master_link

        # Make sure that the PATH environment variable is set
        if not os.environ.get('PATH'):
            os.environ['PATH'] = '/sbin:/usr/sbin:/bin:/usr/bin'

    def list_alternatives(self):
        '''Get the list of alternatives for the master link'''
        dev_null = open('/dev/null', 'w')
        alternatives = []
        p1 = Popen([self.__command, '--list', self.__master_link], stdout=PIPE, stderr=dev_null)
        p = p1.communicate()[0]
        dev_null.close()
        c = p.split('\n')
        for line in c:
            line.strip() and alternatives.append(line.strip())

        return alternatives

    def get_current_alternative(self):
        '''Get the alternative in use'''
        dev_null = open('/dev/null', 'w')
        current_alternative = None
        p1 = Popen([self.__command, '--query', self.__master_link], stdout=PIPE, stderr=dev_null)
        p = p1.communicate()[0]
        dev_null.close()
        c = p.split('\n')
        for line in c:
            if line.strip().startswith('Value:'):
                return line.replace('Value:', '').strip()
        return None

    def get_alternative_by_name(self, name):
        '''Get the alternative link by providing the driver name'''
        alternatives = self.list_alternatives()

        for alternative in alternatives:
            if alternative.__contains__(name):
                return alternative

        return None

    def get_open_drivers_alternative(self):
        '''Get the alternative link for open drivers'''
        return self.get_alternative_by_name(self.__open_drivers_alternative)

    def update_gmenu(self):
        '''Trigger gmenu so that the icons will show up in the menu'''
        try:
            subprocess.check_call(['dpkg-trigger', '--by-package=fakepackage', 'gmenucache'], stdout=open(os.devnull, 'w'))
            subprocess.check_call(['dpkg', '--configure', '-a'], stdout=open(os.devnull, 'w'))
        except (OSError, CalledProcessError):
            pass

    def set_alternative(self, path):
        '''Tries to set an alternative and returns the boolean exit status'''
        try:
            subprocess.check_call([self.__command, '--set', self.__master_link, path], stdout=open(os.devnull, 'w'))
            self.ldconfig()
        except CalledProcessError:
            return False

        self.update_gmenu()

        return True

    def ldconfig(self):
        '''Call ldconfig'''
        try:
            subprocess.check_call(['ldconfig'], stdout=open(os.devnull, 'w'))
        except CalledProcessError:
            return False
        return True


class Switcher(object):

    def __init__(self):
        self.__gl_switcher = Alternatives('gl_conf')

    def __simplify_x_alternative_name(self, alternative):
        return alternative.split('/')[-1]

    def __simplify_gl_alternative_name(self, alternative):
        return alternative.split('/')[-2]

    def __get_gl_alternatives_list(self):
        raw_alternatives = self.__gl_switcher.list_alternatives()
        return  [ self.__simplify_gl_alternative_name(x) for x in raw_alternatives ]

    def __update_initramfs(self):
        subprocess.call(['update-initramfs', '-u'])
        # This may not be necessary
        subprocess.call(['update-initramfs', '-u', '-k', os.uname()[2]])

    def __get_current_alternative(self):
        # This is a list as the 2nd item could be another alternative
        # we might be looking for e.g. an xorg.conf alternative
        alternatives = [None, None]
        raw_gl_alternative = self.__gl_switcher.get_current_alternative()

        if (raw_gl_alternative != None):
            alternatives[0] = self.__simplify_gl_alternative_name(raw_gl_alternative)

        return alternatives

    def enable_alternative(self, alternative):
        # Make sure that the alternative exists
        gl_alternative = self.__gl_switcher.get_alternative_by_name(alternative)

        # See if they are null
        # Abort if gl_alternative is null
        if gl_alternative:
            success = self.__gl_switcher.set_alternative(gl_alternative)
            return success
        else:
            sys.stderr.write("Error: no alternative can be found for %s\n" % alternative)

        return False

    def print_current_alternative(self):
        alternatives = self.__get_current_alternative()

        try:
            alternative = str(alternatives[0])
        except IndexError:
            # No alternative exists
            return False

        if alternative == 'fglrx':
            print 'amd'
        elif alternative == 'pxpress':
            print 'intel'
        else:
            print 'unknown'

        return True

def check_root():
    if not os.geteuid() == 0:
        sys.stderr.write("This operation requires root privileges\n")
        exit(1)

def handle_alternatives_error(mode):
    sys.stderr.write('Error: %s mode can\'t be enabled\n' % mode)
    exit(1)

def handle_query_error():
    sys.stderr.write("Error: no alternative can be found\n")
    exit(1)

def usage():
    sys.stderr.write("Usage: %s amd|intel|query\n" % (sys.argv[0]))

if __name__ == '__main__':
    try:
        arg = sys.argv[1]
    except IndexError:
        arg = None

    check_root()

    if len(sys.argv[1:]) != 1:
        usage()
        exit(1)

    switcher = Switcher()

    if arg == "amd":
        if not switcher.enable_alternative('fglrx'):
            handle_alternatives_error(arg)
    elif arg == "intel":
        if not switcher.enable_alternative('pxpress'):
            handle_alternatives_error(arg)
    elif arg == "query":
        if not switcher.print_current_alternative():
            handle_query_error()
    else:
        usage()
        sys.exit(1)

    exit(0)
