#!/usr/bin/python
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Library General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
#
# Copyright 2003, 2004 Peter Backlund
# Copyright 2004 Thorsten Leemhuis

import ixf86config
import commands
import string
import os
import sys

#
# Adds a module to the Modules section. 
# The module will only be listed once.
#
def addModule(modules, moduleName):
    m = ixf86config.XF86ConfLoad()
    m.name = moduleName
    removeModule(modules, moduleName)
    modules.load.insert(m)

#
# Removes a module from the Modules section.
#
def removeModule(modules, moduleName):
    i = 0
    for m in modules.load:
        if (m.name == moduleName):
            modules.load.remove(i)
        i += 1

#
# Set driver to newDriver where 
# driver is oldDriver
# 
def toggleDriver(device, oldDriver, newDriver):
    for dev in device:
        if (dev.driver == oldDriver):
            dev.driver = newDriver

#
# Add option to driver device section where 
# if dev.driver is driver
# 
def addOption(device, driver, addOptName, addOptVal):
    for dev in device:
        if (dev.driver == driver):
           res = filter (lambda o: o.name == addOptName, dev.options)
           if len(res) == 0:
               dev.options.insert(ixf86config.XF86Option (addOptName,addOptVal))
#
# Remove option to driver device section where 
# if dev.driver is driver
# 
def removeOption(device, driver, addOptName):
    for dev in device:
        if (dev.driver == driver):
           optcount = -1 
           for opt in dev.options:
              optcount += 1
              if opt.name == addOptName:
                 dev.options.remove(optcount)
#
# Create a file in /etc/ld.so.conf.d with a include path
# 
def createLdconfigPathIncludeFile(content):
    try:
        ldIncludeFile=open('/etc/ld.so.conf.d/ati-fglrx.conf', 'w')
        ldIncludeFile.write(content)
        ldIncludeFile.close()
    except:
        None

#
# Remove the file in /etc/ld.so.conf.d with the include path
# 
def removeLdconfigPathIncludeFile():
    try:
        os.remove("/etc/ld.so.conf.d/ati-fglrx.conf")
    except:
        None
  
#
# run ldconfig
# 
def runLdconfig():
    os.system("/sbin/ldconfig")

#
# Create a file /etc/sysconfig/ati-fglrx-config-display with content
# 
def createStatusFile(content):
    try:
        ldIncludeFile=open('/etc/sysconfig/ati-fglrx-config-display', 'w')
        ldIncludeFile.write(content)
        ldIncludeFile.close()
    except:
        None

#
# Checks if a kernel module is loaded and tries to unload if unload = True 
# 
def checkIfKernelModuleIsLoaded(modulename):
    cmd = "/sbin/lsmod | grep " + "\"^" + modulename + "\""
    moduleline = commands.getoutput(cmd)
    if len(moduleline) > 0:
       # module is loaded
       module, size, usage = moduleline.split()
       if usage == "0":
           # usage 0 
           return 1
       else:
           return 2  
    return 0

def checkForModuleAndRemove(modulename):
    #
    # We cannot disable or enable the driver blindly -- the kernel-module might be loaded
    # Try to unload the module. otherwise let the init-script enable or disable the driver during next restart
    #
    ret = checkIfKernelModuleIsLoaded(modulename)
    if ret == 1:
       if os.system("/sbin/modprobe -r "+modulename) > 0 :
           return False
    elif ret == 2:
       return False
    return True


#
# Set Default Color Depth
# 
def setDefaultDepth(device, newdepth):
    for scr in device:
        olddepth = scr.defaultdepth
        if scr.defaultdepth == newdepth:
            continue
        scr.defaultdepth = newdepth
        for dsp in scr.display:
            if dsp.depth == olddepth:
                dsp.depth = newdepth
            elif dsp.depth == newdepth:
                dsp.depth = olddepth


# ------------ 
# Main section
# ------------

# Read config file
(xconfig, xconfigpath) = ixf86config.readConfigFile()

if (len(sys.argv) > 1):
   arg = sys.argv[1]
else:
   print "Usage: ati-fglrx-config-display enable|disable"
   sys.exit(1)

if (len(sys.argv) > 1 and (arg == "enable" or arg == "disable" or arg == "enabledri")):
   #
   # Backup original X config file to .backup-ati-fglrx
   backup_file = None
   output_file = xconfigpath
   
   if os.access(output_file, os.F_OK):
       backup_file = output_file + ".backup-ati-fglrx"
       os.rename(output_file, backup_file)
   else:
       sys.exit(1)

   try:
       if (arg == "enabledri"):
            #
            # Enable DRI
            # When fglrx was loaded we can't enable DRI. The Init-Script knows that and
            # will enable DRI with this parameter on the next system shutdown.
            #
            addModule(xconfig.modules, "dri")
            createStatusFile("FGLRXSTATUS=disabled")
       elif (arg == "enable"):
            #
            # Enable fglrx driver, disable radeon driver
            #
            toggleDriver(xconfig.device, "radeon", "fglrx")
            addOption(xconfig.device, "fglrx","VideoOverlay" ,"on" )
            addModule(xconfig.modules, "dri")
            if os.uname()[4] == "x86_64" :
               createLdconfigPathIncludeFile("/usr/lib/ati-fglrx/\n/usr/lib64/ati-fglrx/\n")
            else:
               createLdconfigPathIncludeFile("/usr/lib/ati-fglrx/\n")
            runLdconfig()
            if not checkForModuleAndRemove("radeon"):
               print "Enabled the fglrx driver. DRI is currently unavailable because the kernel-module"
               print "radeon is already loaded. Please restart the system -- Or end you X-Session, log"
               print "in again, issue the command"
               print "/sbin/rmmod radeon"        
               print "and end your X-Session once again."
            if not checkForModuleAndRemove("drm"):
               print "Enabled the fglrx driver. DRI is currently unavailable because the kernel-module"
               print "drm is already loaded. Please restart the system -- Or end you X-Session, log"
               print "in again, issue the command"
               print "/sbin/rmmod radeon"        
               print "and end your X-Session once again."
            createStatusFile("FGLRXSTATUS=enabled")
       else:
            #
            # Disable fglrx driver, enable radeon driver
            #
            removeOption(xconfig.device, "fglrx","VideoOverlay")
            toggleDriver(xconfig.device, "fglrx", "radeon")
            removeLdconfigPathIncludeFile()
            runLdconfig()
            if checkIfKernelModuleIsLoaded("fglrx") >= 1 :
               # If fglrx kernel-module is or was loaded the system max restart
               # when we now run the radeon driver with dri
               # so remove the dri module and let the init-script
               # re-enable it on the next shutdown 
               removeModule(xconfig.modules, "dri")
               createStatusFile("FGLRXSTATUS=dritobeenabled")
               print "Disabled the fglrx driver but did not enable DRI because many systems will"
               print "suddenly reboot in this case. DRI will be enabled after system restart."
            else:
               addModule(xconfig.modules, "dri")
               createStatusFile("FGLRXSTATUS=disabled")
       # Write new XF86Config
       xconfig.write(output_file)
   except:
       # Something went wrong, restore backup
       os.rename(backup_file, output_file)
else:
    print "Usage: ati-fglrx-config-display enable|disable|enabledri"
    sys.exit(1)

