#!/usr/bin/python
#  Copyright (c) 2010 Canonical
#
#  Author: Oliver Grawert <ogra@canonical.com>
#
#  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 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
#
######################################################################

import pygtk
pygtk.require('2.0')
import gtk
from subprocess import Popen,PIPE,call,STDOUT

import gettext
from gettext import gettext as _

import os
import fcntl
import shutil

releases = ['lucid', 'karmic', 'jaunty']
icon_theme = gtk.icon_theme_get_default()

try:
    tmpdir = os.environ['TMPDIR']
except:
    tmpdir = '/tmp'

os.chdir(tmpdir)

class MainWin:
    def __init__(self):
        APP="rootstock-gtk"
        DIR="/usr/share/locale"

        gettext.bindtextdomain(APP, DIR)
        gettext.textdomain(APP)

        self.tasks = ""
        self.window = gtk.Window(gtk.WINDOW_TOPLEVEL)
        self.window.connect("destroy", self.destroy)
        self.window.set_default_size(400, 380)
        self.window.set_title(_('Build Root Filesystem'))

        self.outdir = os.environ['HOME']

        try:
            pixbuf = icon_theme.load_icon("rootstock_64", 24, 0)
        except:
            pass
        gtk.window_set_default_icon_list(pixbuf)

        # the menubar
        mb = gtk.MenuBar()
        filemenu = gtk.Menu()
        filem = gtk.MenuItem(_('File'))
        filem.set_submenu(filemenu)

        gtk.stock_add([(gtk.STOCK_DIRECTORY, _('Set Target Directory'), 0, 0, "")])
        sel_dir = gtk.ImageMenuItem(gtk.STOCK_DIRECTORY, None)
        sel_dir.connect("activate", self.select_dir)
        filemenu.append(sel_dir)

        gtk.stock_add([(gtk.STOCK_OPEN, _('Load Profile'), 0, 0, "")])
        open_prof = gtk.ImageMenuItem(gtk.STOCK_OPEN, None)
        open_prof.connect("activate", self.open_profile)
        #filemenu.append(open_prof)

        gtk.stock_add([(gtk.STOCK_SAVE, _('Save Profile'), 0, 0, "")])
        save_prof = gtk.ImageMenuItem(gtk.STOCK_SAVE, None)
        save_prof.connect("activate", self.save_profile)
        #filemenu.append(save_prof)

        sep = gtk.SeparatorMenuItem()
        filemenu.append(sep)

        exit = gtk.ImageMenuItem(gtk.STOCK_QUIT, None)
        exit.connect("activate", gtk.main_quit)
        filemenu.append(exit)

        mb.append(filem)

        # topframe
        tf = gtk.Frame()
        tflbl = gtk.Label()
        tf.set_label_widget(tflbl)
        tflbl.set_markup(_('<b>General Options</b>'))

        top_table = gtk.Table(5, 2, False)
        top_hbox = gtk.HBox(True, 0)
        top_vbox = gtk.VBox(True, 0)
        top_hbox.pack_start(top_vbox, False, True, 5)
        top_vbox.pack_start(top_table, False, True, 5)
        tf.add(top_hbox)

        mirrorlabel = gtk.Label(_('Mirror'))
        mirrorlabel.set_alignment(xalign=0, yalign=0.5)
        self.mirrorentry = gtk.Entry()
        self.mirrorentry.insert_text("http://ports.ubuntu.com/ubuntu-ports", 0)
        self.mirrorentry.set_position(-1)
        top_table.attach(mirrorlabel, 0, 1, 0, 1)
        top_table.attach(self.mirrorentry, 1, 2, 0, 1)

        self.xmirror = gtk.CheckButton(_('Use additional mirror'))
        self.xmirrorentry = gtk.Entry()
        self.xmirrorentry.set_sensitive(False)
        self.xmirror.connect("toggled", self.set_sensitivity)
        top_table.attach(self.xmirror, 0, 1, 1, 2)
        top_table.attach(self.xmirrorentry, 1, 2, 1, 2)

        releaselabel = gtk.Label(_('Release'))
        releaselabel.set_alignment(xalign=0, yalign=0.5)
        self.releaseselect = gtk.combo_box_new_text()
        for release in releases:
            self.releaseselect.append_text(release)
        self.releaseselect.set_active(0)
        top_table.attach(releaselabel, 0, 1, 2, 3)
        top_table.attach(self.releaseselect, 1, 2, 2, 3)

        scriptlabel = gtk.Label(_('Custom post installation script'))
        scriptlabel.set_alignment(xalign=0, yalign=0.5)
        gtk.stock_add([(gtk.STOCK_OPEN, _('Load Script'), 0, 0, "")])
        self.scriptfile = gtk.FileChooserButton(_('Select script file'), None)
        top_table.attach(scriptlabel, 0, 1, 3, 4)
        top_table.attach(self.scriptfile, 1, 2, 3, 4)

        imgsizelabel = gtk.Label(_('Rootfs build Space'))
        imgsizelabel.set_alignment(xalign=0, yalign=0.5)
        imgsizebox = gtk.HBox(False,0)
        self.imgsize = gtk.Adjustment(4, 1, 100, 1, 0, 0)
        imgsizespin = gtk.SpinButton(self.imgsize, 1, 0)
        imgsizeunit = gtk.Label("Gigabyte")
        imgsizeunit.set_alignment(xalign=0, yalign=0.5)
        imgsizebox.add(imgsizespin)
        imgsizebox.add(imgsizeunit)
        top_table.attach(imgsizelabel, 0, 1, 4, 5)
        top_table.attach(imgsizebox, 1, 2, 4, 5)

        # middleframe
        mf = gtk.Frame()
        mflbl = gtk.Label()
        mf.set_label_widget(mflbl)
        mflbl.set_markup(_('<b>Options for packages to install</b>'))

        middle_table = gtk.Table(3, 2, False)
        mid_hbox = gtk.HBox(True, 0)
        mid_vbox = gtk.VBox(True, 0)
        mid_hbox.pack_start(mid_vbox, False, True, 5)
        mid_vbox.pack_start(middle_table, False, True, 5)
        mf.add(mid_hbox)

        pkglistlabel = gtk.Label(_('Use a list of packages'))
        pkglistlabel.set_alignment(xalign=0, yalign=0.5)
        self.pkglistentry = gtk.Entry()
        self.pkglistentry.insert_text("ubuntu-minimal", 0)
        middle_table.attach(pkglistlabel, 0, 1, 1, 2)
        middle_table.attach(self.pkglistentry, 1, 2, 1, 2)

        tasklabel = gtk.Label(_('Select an Ubuntu task'))
        tasklabel.set_alignment(xalign=0, yalign=0.5)
        tasksel = gtk.Button(_('Run Tasksel'))
        tasksel.connect("clicked", self.run_tasksel)
        middle_table.attach(tasklabel, 0, 1, 2, 3)
        middle_table.attach(tasksel, 1, 2, 2, 3)

        seedfilelabel = gtk.Label(_('Use a seed file'))
        seedfilelabel.set_alignment(xalign=0, yalign=0.5)
        gtk.stock_add([(gtk.STOCK_OPEN, _('Load Seed File'), 0, 0, "")])
        self.seedfile = gtk.FileChooserButton(_('Select seed file'), None)
        middle_table.attach(seedfilelabel, 0, 1, 3, 4)
        middle_table.attach(self.seedfile, 1, 2, 3, 4)

        # bottomframe
        bf = gtk.Frame()
        bflbl = gtk.Label()
        bf.set_label_widget(bflbl)
        bflbl.set_markup(_('<b>Advanced Options</b>'))

        bottom_table = gtk.Table(2, 2, False)
        bot_hbox = gtk.HBox(True, 0)
        bot_vbox = gtk.VBox(True, 0)
        bot_hbox.pack_start(bot_vbox, False, True, 5)
        bot_vbox.pack_start(bottom_table, False, True, 5)
        bf.add(bot_hbox)

        seriallabel = gtk.Label(_('Use serial device in rootfs'))
        seriallabel.set_alignment(xalign=0, yalign=0.5)
        self.serialentry = gtk.Entry()
        bottom_table.attach(seriallabel, 0, 1, 1, 2)
        bottom_table.attach(self.serialentry, 1, 2, 1, 2)

        #partitionlabel = gtk.Label(_('Partition to unpack rootfs to'))
        #partitionlabel.set_alignment(xalign=0, yalign=0.5)
        #self.partitionentry = gtk.Entry()
        #bottom_table.attach(partitionlabel, 0, 1, 2, 3)
        #bottom_table.attach(self.partitionentry, 1, 2, 2, 3)

        # buttonbar
        bb = gtk.HBox(False, 0)
        bbhbox = gtk.HBox(False, 0)

        gtk.stock_add([(gtk.STOCK_EXECUTE, _('Build'), 0, 0, "")])
        button1 = gtk.Button(None, stock=gtk.STOCK_EXECUTE)
        button1.connect("clicked", self.run_build)

        button2 = gtk.Button(None, stock=gtk.STOCK_CANCEL)
        button2.connect("clicked", gtk.main_quit)

        bb.pack_end(button1, False, False, 0)
        bb.pack_end(button2, False, False, 0)
        bbhbox.pack_end(bb, False, False, 5)

        # toplevel layout
        mainbox = gtk.VBox(False, 0)
        mainboxv = gtk.VBox(False, 0)
        mainboxh = gtk.HBox(True, 0)

        mainbox.pack_start(tf, False, False, 5)
        mainbox.pack_start(mf, False, False, 5)
        mainbox.pack_start(bf, False, False, 5)

        mainboxh.pack_start(mainbox, False, True, 5)

        mainboxv.pack_start(mb, False, False, 0)
        mainboxv.pack_start(mainboxh, False, True, 5)
        mainboxv.pack_end(bbhbox, False, False, 5)

        self.window.add(mainboxv)
        self.window.show_all()

    def set_sensitivity(self, widget):
        self.xmirrorentry.set_text("")
        self.xmirrorentry.set_sensitive(self.xmirror.get_active())

    def run_build(self, widget):
        command = ['rootstock -f rootstockfs']
        if self.mirrorentry.get_text():
            command.append(' -m '+self.mirrorentry.get_text())
        if self.xmirrorentry.get_text():
            command.append(' --extra-mirror '+self.xmirrorentry.get_text())
        command.append(' -d '+releases[self.releaseselect.get_active()])
        if self.scriptfile.get_filename() != None:
            command.append(' --script '+str(self.scriptfile.get_filename()))
        command.append(' -i '+str(int(self.imgsize.get_value()))+'G')
        if self.pkglistentry.get_text():
            command.append(' -s '+self.pkglistentry.get_text()+' '+str(self.tasks))
        if self.seedfile.get_filename() != None:
            command.append(' --manifest '+str(self.seedfile.get_filename()))
        if self.serialentry.get_text():
            command.append(' --serial '+self.serialentry.get_text())
        while gtk.events_pending():
            gtk.main_iteration(True)
        #print 'Target partition: '+self.partitionentry.get_text()
        self.draw_progress_win()
        self.run_command(command)

    def draw_progress_win(self):
        while gtk.events_pending():
            gtk.main_iteration(True)
        self.window.set_sensitive(False)
        self.prgwin = gtk.Window(gtk.WINDOW_TOPLEVEL)
        self.prgwin.set_transient_for(self.window)
        self.prgwin.connect("destroy", self.close_progress)
        self.prgwin.set_title(_('Building Root Filesystem'))

        self.prgwin.set_default_size(300, 80)

        hbox = gtk.HBox(False, 5)
        framebox = gtk.HBox(False, 5)
        vbox = gtk.VBox(False, 5)
        progressheadbox = gtk.HBox(False, 5)

        headericon = gtk.Image()
        headericon.set_from_stock(gtk.STOCK_EXECUTE, gtk.ICON_SIZE_DIALOG)
        progressheader = gtk.Label()
        progressheader.set_markup(_('<b>Building Root Filesystem...</b>\n This can take a while, please be patient'))
        progressheader.set_alignment(xalign=0.0, yalign=0.5)

        progressheadbox.pack_start(headericon, False, False, 5)
        progressheadbox.pack_start(progressheader, False, False, 0)

        self.progress = gtk.ProgressBar()

        self.progresslabel = gtk.Label()
        self.progresslabel.set_alignment(xalign=0.0, yalign=0.5)
        sep = gtk.HSeparator()

        vbox.pack_start(progressheadbox, True, True, 5)
        vbox.pack_start(self.progress, False, False, 0)
        vbox.pack_start(self.progresslabel, True, True, 0)
        vbox.pack_start(sep, True, True, 0)

        self.closebut = gtk.Button(None, stock=gtk.STOCK_CANCEL)
        self.closebut.connect("clicked", self.close_progress)
        hbox.pack_end(self.closebut, False, False, 5)

        vbox.pack_end(hbox, False, False, 5)
        framebox.pack_start(vbox, True, True, 5)

        self.prgwin.add(framebox)
        self.prgwin.show_all()

    def run_command(self, command):
        cmdstr = ""
        lines = 1.0
        tarball = ''
        for item in command:
            cmdstr=cmdstr+' '+item
        output = Popen([cmdstr], stdout=PIPE, stderr=None, shell=True)
        flags = fcntl.fcntl(output.stdout, fcntl.F_GETFL)
        fcntl.fcntl(output.stdout, fcntl.F_SETFL, flags | os.O_NONBLOCK)

        self.rootstockpid = output.pid

        self.progresslabel.set_text(_('Preparing environment for rootfs build ...'))

        while output.poll() is None:
            while gtk.events_pending():
                gtk.main_iteration(True)
            try:
                line = output.stdout.readline()
                print line.strip()
                if line.startswith('I: Running first stage'):
                    ptext = _('Running first stage of build process ...')
                    lines = 1.0
                    frac = 850.0
                if line.startswith('I: Switching to Virtual Machine'):
                    ptext = _('Configuring virtual machine for second stage processing ...')
                    lines = 1.0
                    frac = 40.0
                if line.startswith('packagecount='):
                    lines = 1.0
                    ptext = _('Running second stage package install in virtual machine ...')
                    pkgcount = float(line.split('=')[1])
                    if pkgcount > 0.0:
                        frac = (pkgcount * 4.0) + 500
                    else:
                        frac = 1.0
                if line.startswith('I: Enabling firstboot configuration'):
                    ptext = _('Enabling first boot user configuration assistant ...')
                    lines = 1.0
                    frac = 350.0
                if line.startswith('I: First stage install done'):
                    lines = 850.0
                if line.startswith('I: Virtual Machine done'):
                    lines = 350.0
                if line.startswith('I: Creating tarball from rootfs'):
                    ptext = _('Creating final tarball, please wait ...')
                    lines = 1.0
                    frac = 2
                if line.startswith('I: Cleaning up...'):
                    ptext = _('Done, cleaning up build environment ...')
                    lines = 1.0
                    frac = 1
                prog_count = float(lines/frac)
                if prog_count > 1.0:
                    prog_count = 1.0
                self.progress.set_fraction(prog_count)
                self.progresslabel.set_text(ptext)
                lines += 1
                if line.startswith('I: ARM rootfs created'):
                    tarball = line.split('as ')[1].split('/')[-1].strip()
                if line.startswith('I: A logfile was saved'):
                    self.logfile = line.split(' as ')[1].split('/')[-1].strip()
            except: pass
        self.prgwin.hide()
        self.window.set_sensitive(True)
        retval = output.returncode
        if tarball != '':
            self.success(tarball)
        elif retval != 0:
            self.error(_('Something went wrong, rootstock died with error code: %i') % retval)

    def close_progress(self, widget):
        while gtk.events_pending():
            gtk.main_iteration(True)
        if self.rootstockpid+1:
            try: os.kill(self.rootstockpid+1, 6)
            except: pass
            widget.get_parent_window().hide()
            self.window.set_sensitive(True)
        else:
            widget.get_parent_window().hide()
            self.window.set_sensitive(True)

    def error(self, string):
        string = _('<b>Error:</b>\n %s') % string
        try:
            if self.logfile:
                if tmpdir != self.outdir:
                    shutil.move(tmpdir+'/'+self.logfile, self.outdir)
                logfile = '%s/%s' % (self.outdir, self.logfile)
                logstr = _(' A logfile was saved in %s\n') % logfile
                string = '%s\n%s' % (string, logstr)
        except:
            pass
        dialog = gtk.Dialog(_('Error'),
                self.window,
                gtk.DIALOG_MODAL | gtk.DIALOG_DESTROY_WITH_PARENT,
                (gtk.STOCK_OK, gtk.RESPONSE_ACCEPT))
        hbox = gtk.HBox(False, 5)
        image = gtk.image_new_from_stock(gtk.STOCK_DIALOG_ERROR, gtk.ICON_SIZE_DIALOG)
        label = gtk.Label()
        label.set_markup(string)
        hbox.pack_start(image, False, False, 0)
        hbox.pack_start(label, False, False, 5)
        dialog.vbox.pack_start(hbox)
        dialog.show_all()
        result = dialog.run()
        if result:
            gtk.main_quit()

    def success(self, tarball):
        if tmpdir != self.outdir:
            shutil.move(tmpdir+'/'+tarball, self.outdir)
            try:
                shutil.move(tmpdir+'/'+self.logfile, self.outdir)
            except:
                pass
        header = _('<b>Build successful:</b>\n')
        tarstr = _('The filesystem archive was created in %s as:\n <b>%s</b>\n\n') % (self.outdir, tarball)
        string = _('%s %s The logfile of this build, %s\n was saved in the same directory.') % (header, tarstr, self.logfile)
        dialog = gtk.Dialog(_('Successfully Built Root Filesystem'),
                self.window,
                gtk.DIALOG_MODAL | gtk.DIALOG_DESTROY_WITH_PARENT,
                (gtk.STOCK_OK, gtk.RESPONSE_ACCEPT))
        hbox = gtk.HBox(False, 5)
        image = gtk.image_new_from_stock(gtk.STOCK_DIALOG_INFO, gtk.ICON_SIZE_DIALOG)
        label = gtk.Label()
        label.set_markup(string)
        hbox.pack_start(image, False, False, 0)
        hbox.pack_start(label, False, False, 5)
        dialog.vbox.pack_start(hbox)
        dialog.show_all()
        result = dialog.run()
        if result:
            gtk.main_quit()

    def run_tasksel(self, widget):
        self.tasks = Popen(["/usr/lib/rootstock/rootstock-tasksel"], stdout=PIPE).communicate()[0].rstrip()

    def select_dir(self, widget):
        chooser = gtk.FileChooserDialog(_('Select Target Directory'), self.window, gtk.FILE_CHOOSER_ACTION_SELECT_FOLDER, (gtk.STOCK_OK, gtk.RESPONSE_OK), None)
        chooser.set_current_folder_uri('file://'+self.outdir)
        if chooser.run():
            self.outdir = chooser.get_current_folder_uri().split('file://')[1]
            chooser.destroy()

    def open_profile(self, widget):
        print 'open'

    def save_profile(self, widget):
        print 'save'

    def destroy(self, widget, data=None):
        gtk.main_quit()

    def main(self):
        gtk.main()

if __name__ == "__main__":
    app = MainWin()
    app.main()
