#!/usr/bin/python

import argparse
import sys
import os
import subprocess
import pyinotify
import time
import gtk
import Image
import atexit

parser = argparse.ArgumentParser()
parser.add_argument('--qmp-shell-root', default='/home/alon/src/spice_upstream/qemu/QMP')
parser.add_argument('--host', default='localhost')
parser.add_argument('--port', default=0, type=int)
parser.add_argument('--path', default='')
parser.add_argument('--filename', default='/tmp/a.ppm')
parser.add_argument('--show', default='showimage')
parser.add_argument('--noerase', dest='erase', default=True, action='store_false')
parser.add_argument('--id', default='')
parser.add_argument('--times', default=1, type=int)
parser.add_argument('--sleep', default=0, type=int)
args = parser.parse_args(sys.argv[1:])
if not args.path and not args.port:
    qmp_addr = '/tmp/qmp'
elif args.path:
    qmp_addr = args.path
else:
    qmp_addr = (args.host, args.port)
print "connecting to %s" % str(qmp_addr)
qmp_shell = '%s/qmp-shell' % args.qmp_shell_root
if not os.path.exists(qmp_shell):
    print("qmp_shell not at %s" % qmp_shell)
    os.exit(1)
sys.path.append(args.qmp_shell_root)
import qmp
mon = qmp.QEMUMonitorProtocol(address=qmp_addr)
mon.connect()

def get_notifier_events():
    notifier.process_events()
    if notifier.check_events(10):
        # read notified events and enqeue them
        notifier.read_events()
    return True

class GUI:
    def __init__(self, args):
        self.window = gtk.Window()
        self.image = gtk.Image()
        self.vbox = gtk.VBox()
        self.window.add(self.vbox)
        self.vbox.add(self.image)
        self.button = gtk.Button()
        self.button.set_label('take screenshot')
        self.button.connect("clicked", self.take_screenshot)
        self.vbox.add(self.button)
        self.current_image = 0
        self.times = args.times
        self.args = args
        self.window.show_all()

    def take_screenshot(self, _button):
        # use timeout_add to avoid being inside the button callback.
        gtk.timeout_add(0, self.take_screen_shot)

    def show_image(self, filename):
        im = Image.open('/tmp/a.ppm')
        buf=gtk.gdk.pixbuf_new_from_data(im.tostring(),
                gtk.gdk.COLORSPACE_RGB, 0, 8, im.size[0], im.size[1],
                im.size[0]*3)
        self.image.set_from_pixbuf(buf)
        self.current_image += 1
        self.window.set_title(str(self.current_image))
        # proceed in a different stack to allow update of screen
        gtk.timeout_add(args.sleep, self.take_screen_shot)

    def take_screen_shot(self):
        if self.current_image < self.times:
            take_screen_dump(self.args.id, self.args.filename)
        else:
            raw_input('press enter to quit')
            sys.exit(0)
        return False

gui = GUI(args)

if os.path.exists(args.filename):
    print "removing old %s" % args.filename
    os.unlink(args.filename)

def cleanup():
    if args.erase:
        print("erasing %s" % args.filename)
        os.unlink(args.filename)

atexit.register(cleanup)

# watch for write to target
class ProcessFile(pyinotify.ProcessEvent):
    def process_IN_CLOSE_WRITE(self, event):
        if event.name != os.path.basename(args.filename):
            return
        gui.show_image(args.filename)

if args.show:
    wm = pyinotify.WatchManager()
    notifier = pyinotify.Notifier(wm, ProcessFile())
    wdd = wm.add_watch(os.path.dirname(args.filename), pyinotify.IN_CLOSE_WRITE)

def take_screen_dump(qxl_id, filename):
    print "saving qxl_id %r to %s" % (qxl_id, filename)
    if qxl_id != '':
        mon.cmd('__com.redhat_qxl_screendump', dict(id=str(qxl_id),
                                filename=filename))
    else:
        mon.cmd('screendump', dict(filename=filename))

take_screen_dump(args.id, args.filename)

# ugly - can't we connect gtk to inotify directly
gtk.timeout_add(10, get_notifier_events)
gtk.mainloop()