#!/usr/bin/python

# This file is /etc/pm/sleep.d/00compiz

import os
import sys
import time
import cPickle as pickle
import tempfile

def psinfo(pid, field):
    return os.popen('ps --no-headers -ww -o %s %s' % (field, pid)).read()[:-1]

def psenviron(pid):
    return dict(
        x.split('=',1)
        for x in open('/proc/%s/environ' % pid).read().split('\0')
        if '=' in x)

def save_environ(pids):

    # Record information about the processes that invoked compiz.real so that we
    # can re-launch them upon resume.
    fd, picklefile = tempfile.mkstemp('.pck', 'compiz-fglrx', '/var/run')
    
    try:
        parent_pids=[int(x) for x in pids]

        parent_info = dict([
            (p, (psenviron(p), psinfo(p, 'cmd'), psinfo(p, 'user')))
             for p in parent_pids]
            )

        pickle.dump(parent_info, open(picklefile, 'w'),2)

        # print the pickle file name on stdout so the calling shell script
        # can record it in an environment variable.  Merely setting a value
        # in os.environ here would not cause the variable to be exported.
        print picklefile
    except:
        os.close(fd)
    
def restore_environ():
    console = os.environ['COMPIZ_CONSOLE']
    picklefile = os.environ['COMPIZ_PARENT_INFO']

    parent_info = pickle.load(open(picklefile))
    os.unlink(picklefile)

    # == superceded by https://bugs.launchpad.net/ubuntu/+source/linux/+bug/218760/comments/21 ==
    # This is really a fix for a different bug, which causes resume to fail
    # when you suspend while docked with USB devices in the dock.  If you'll
    # never use a dock, you can eliminate these two lines and avoid some
    # flashing.  One day I'll factor this out but right now it's working and
    # I don't want to break anything.
    # os.system('/usr/bin/chvt %s' % console)
    # os.system('/usr/bin/chvt 63')

    for pid, (env,cmd,user) in parent_info.items():

        if 'DISPLAY' in env and 'XAUTHORITY' in env and os.path.isfile(env['XAUTHORITY']):

            try:
                os.kill(pid,15)
            except OSError: # It may have exited when compiz.real was killed
                pass 

            save_env=os.environ.copy()
            try:
                os.environ.clear()
                os.environ.update(env)
                os.system('sudo -E -b -u "%s" %s' % (user,cmd))
            finally:
                os.environ.clear()
                os.environ.update(save_env)

    # Without this we end up in some VT limbo land and have to press
    # ctrl-alt-f7 to get back
    os.system('/usr/bin/chvt %s' % console)

if __name__ == '__main__':

    # The standard output of the save_environ command will be captured by,
    # /etc/pm/config.d/00-compiz-stop.sh, so if you need to debug this script
    # you should print to sys.stderr, e.g.
    #
    #   print >>sys.stderr, sys.argv
    
    try:
        if sys.argv[1] == 'save_environ':
            save_environ(sys.argv[2:])
            
        elif sys.argv[1] in ('resume','thaw'):
            restore_environ()

        else: # suspend,hibernate: nothing to do
            pass

    except Exception, e:
        print '00compiz:', str(e)
        raise
    
exit(0)