#!/usr/bin/env python

import time

import gobject, gtk, pango, cairo

import gtksourceview2 as gtksourceview

import psycopg2

import psycopg2.extensions

import psycopg2.extras

from pgconsole.config import cfg

from pgconsole.editor import Editor

from pgconsole.dataview import DataView

from pgconsole.database import Database, BadConnectionError, DatabaseError

from pgconsole.settings import Settings

from pgconsole.panedext import HPanedExt, VPanedExt

class PgConsoleApp:

    def __init__(self):

        self.db = Database()

        self.conn = None

        win = gtk.Window(gtk.WINDOW_TOPLEVEL)

        self.win = win

        win.set_title('PostgreSQL Console')

        win.set_size_request(300, 200)  # minimal size

        win.set_default_size(800, 600)

        self.restore_window_size()

        win.connect("destroy", self.destroy)

        win.connect("key_press_event", self.keypress)

        win.connect('configure-event', self.on_configure)

        # toolbar

        toolbar = gtk.Toolbar()

        toolbar.set_style(gtk.TOOLBAR_ICONS)

        toolbar.set_property("icon-size", gtk.ICON_SIZE_SMALL_TOOLBAR)

        tb = gtk.ToolButton(gtk.STOCK_PREFERENCES)

        tb.set_label('Settings')

        tb.set_tooltip_text('Settings')

        tb.connect('clicked', self.settings)

        self.tb_settings = tb

        toolbar.add(tb)

        self.cb_server = gtk.combo_box_entry_new_text()

        self.cb_server.set_tooltip_text('Server')

        self.cb_server.get_child().set_property("editable", False)

        self.cb_server.set_property("add_tearoffs", True)

        #self.cb_server.set_property("focus-on-click", False)

        self.cb_server.set_property("can-focus", True)

        self.cb_server.connect('changed', self.on_server_changed)

        self.cb_server.connect("key_press_event", self.toolbar_server_keypress)

        ti = gtk.ToolItem()

        ti.add(self.cb_server)

        toolbar.add(ti)

        self.cb_dbname = gtk.combo_box_entry_new_text()

        self.cb_dbname.set_tooltip_text('Database')

        self.cb_dbname.get_child().set_property("editable", False)

        self.cb_dbname.set_property("add_tearoffs", True)

        #self.cb_dbname.set_property("focus-on-click", False)

        self.cb_dbname.set_property("can-focus", True)

        ti = gtk.ToolItem()

        ti.add(self.cb_dbname)

        toolbar.add(ti)

        tb = gtk.ToolButton(gtk.STOCK_CONNECT)

        tb.set_label('Connect')

        tb.set_tooltip_text('Connect')

        tb.connect('clicked', self.connect)

        toolbar.add(tb)

        self.tb_connect = tb

        sep = gtk.SeparatorToolItem()

        toolbar.add(sep)

        tb = gtk.ToolButton(gtk.STOCK_EXECUTE)

        tb.set_label('Execute')

        tb.set_tooltip_text('Execute')

        tb.connect('clicked', self.execute)

        toolbar.add(tb)

        tb.set_sensitive(False)

        self.tb_execute = tb

        tb = gtk.ToolButton(gtk.STOCK_NEW)

        tb.set_label('Begin transaction')

        tb.set_tooltip_text('Begin transaction')

        tb.connect('clicked', self.begin)

        tb.set_sensitive(False)

        toolbar.add(tb)

        self.tb_begin = tb

        tb = gtk.ToolButton(gtk.STOCK_APPLY)

        tb.set_label('Commit')

        tb.set_tooltip_text('Commit')

        tb.connect('clicked', self.commit)

        tb.set_sensitive(False)

        self.tb_commit = tb

        toolbar.add(tb)

        tb = gtk.ToolButton(gtk.STOCK_CANCEL)

        tb.set_label('Rollback')

        tb.set_tooltip_text('Rollback')

        tb.connect('clicked', self.rollback)

        tb.set_sensitive(False)

        self.tb_rollback = tb

        toolbar.add(tb)

        sep = gtk.SeparatorToolItem()

        toolbar.add(sep)

        # editor

        self.editor = Editor()

        # data view

        self.dataview = DataView()

        vbox = gtk.VBox(False, 2)

        sep = gtk.SeparatorToolItem()

        sep.set_expand(True)

        sep.set_draw(False)

        toolbar.add(sep)

        self.throbber_anim = gtk.gdk.PixbufAnimation('data/Throbber.gif')

        self.throbber = gtk.Image()

        tb = gtk.ToolItem()

        tb.add(self.throbber)

        toolbar.add(tb)

        sep = gtk.SeparatorToolItem()

        sep.set_draw(False)

        toolbar.add(sep)

        vbox.pack_start(toolbar, False, False, 0)

        vpaned = VPanedExt()

        vpaned.set_border_width(5)

        hpaned = HPanedExt()

        hpaned.set_border_width(0)

        hpaned.add1(self.editor)

        hpaned.child_set_property(self.editor, 'shrink', False)

        hpaned.set_snap2(80)

        hpaned.set_property('position', 500)

        vpaned.add1(hpaned)

        vpaned.set_snap1(80)

        vpaned.set_property('position', 300)

        vpaned.add2(self.dataview)

        vpaned.set_snap2(80)

        self.vpaned = vpaned

        self.hpaned = hpaned

        # log

        self.logbuf = gtk.TextBuffer()

        view = gtk.TextView(self.logbuf)

        view.set_editable(False)

        font_desc = pango.FontDescription('monospace')

        if font_desc:

            view.modify_font(font_desc)

        sw = gtk.ScrolledWindow()

        sw.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)

        sw.set_shadow_type(gtk.SHADOW_ETCHED_IN)

        sw.add(view)

        hpaned.add2(sw)

        vbox.pack_start(vpaned, padding=0)

        win.add(vbox)

        self.editor.view.grab_focus()

        self.reload_server_list(0)

        win.show_all()

        self.restore_win_state()

    def main(self):

        gtk.main()

    def destroy(self, widget, data=None):

        self.save_win_state()

        cfg.save()

        gtk.main_quit()

    def on_configure(self, w, ev):

        cfg.root.window.size.width = ev.width

        cfg.root.window.size.height = ev.height

    def restore_window_size(self):

        self.win.resize(cfg.root.window.size.width, cfg.root.window.size.height)

    def save_win_state(self):

        cfg.root.window.dividers.verticaldivider = self.vpaned.get_position()

        cfg.root.window.dividers.horizontaldivider = self.hpaned.get_position()

        cfg.root.window.dividers.editordivider = self.editor.get_position()

    def restore_win_state(self):

        pos = cfg.root.window.dividers.verticaldivider

        if pos >= 0:

            self.vpaned.set_position(pos)

        pos = cfg.root.window.dividers.horizontaldivider

        if pos >= 0:

            self.hpaned.set_position(pos)

        pos = cfg.root.window.dividers.editordivider

        if pos >= 0:

            self.editor.set_position(pos)

    def get_typename(self, oid, size):

        self.curs.execute('SELECT typname FROM pg_type WHERE oid=%s', [oid])

        psycopg2.extras.wait_select(self.curs.connection)

        row = self.curs.fetchone()

        typname = None

        if row:

            typname = row[0]

            if typname == 'int4':

                return 'integer'

            if typname in ('timestamp', 'interval', 'date'):

                return typname

            if size and size > 0 and size < 65535:

                typname += '(%s)' % size

        return typname

    def get_conninfo(self, nodb=False):

        sel = self.cb_server.get_active()

        srv = cfg.servers.server[sel]

        if nodb:

            dbname = 'postgres'

        else:

            dbname = self.cb_dbname.get_active_text()

        conninfo = 'host=%s port=%s dbname=%s user=%s password=%s' \

            % (srv.host, srv.port, dbname, srv.user, srv.password)

        return conninfo

    def connect(self, w):

        conninfo = self.get_conninfo()

        if self.conn:

            # disconnect

            self.db.put_conn(conninfo, self.conn)

            self.conn = None

            self.cb_server.set_sensitive(True)

            self.cb_dbname.set_sensitive(True)

            self.tb_connect.set_stock_id(gtk.STOCK_CONNECT)

            self.tb_connect.set_label('Connect')

            self.tb_connect.set_tooltip_text('Connect')

            self.tb_execute.set_sensitive(False)

            self.tb_begin.set_sensitive(False)

        else:

            # connect

            self.logbuf.insert(self.logbuf.get_end_iter(), 'Connect %s\n' % conninfo)

            try:

                self.conn = self.db.get_conn(conninfo)

            except DatabaseError, e:

                self.logbuf.insert(self.logbuf.get_end_iter(), 'Error:\n%s\n' % e)

                return

            self.cb_server.set_sensitive(False)

            self.cb_dbname.set_sensitive(False)

            self.tb_connect.set_stock_id(gtk.STOCK_DISCONNECT)

            self.tb_connect.set_label('Disconnect')

            self.tb_connect.set_tooltip_text('Disconnect')

            self.tb_execute.set_sensitive(True)

            self.tb_begin.set_sensitive(True)

    def begin(self, w):

        self.logbuf.insert(self.logbuf.get_end_iter(), 'Begin transaction\n')

        curs = self.conn.cursor()

        curs.execute('BEGIN')

        psycopg2.extras.wait_select(curs.connection)

        self.tb_connect.set_sensitive(False)

        self.tb_begin.set_sensitive(False)

        self.tb_commit.set_sensitive(True)

        self.tb_rollback.set_sensitive(True)

    def commit(self, w):

        self.logbuf.insert(self.logbuf.get_end_iter(), 'Commit\n')

        curs = self.conn.cursor()

        curs.execute('COMMIT')

        psycopg2.extras.wait_select(curs.connection)

        self.tb_connect.set_sensitive(True)

        self.tb_begin.set_sensitive(True)

        self.tb_commit.set_sensitive(False)

        self.tb_rollback.set_sensitive(False)

    def rollback(self, w):

        self.logbuf.insert(self.logbuf.get_end_iter(), 'Rollback\n')

        curs = self.conn.cursor()

        curs.execute('ROLLBACK')

        psycopg2.extras.wait_select(curs.connection)

        self.tb_connect.set_sensitive(True)

        self.tb_begin.set_sensitive(True)

        self.tb_commit.set_sensitive(False)

        self.tb_rollback.set_sensitive(False)

    def execute(self, widget):

        query = self.editor.get_selection() or self.editor.get_text()

        self.tb_connect.set_sensitive(False)

        self.tb_execute.set_sensitive(False)

        self.tb_begin.set_sensitive(False)

        self.tb_commit.set_sensitive(False)

        self.tb_rollback.set_sensitive(False)

        self.throbber.set_from_animation(self.throbber_anim)

        self.curs = self.conn.cursor()

        self.t1 = time.time()

        try:

            self.curs.execute(query)

        except (psycopg2.OperationalError, psycopg2.DatabaseError), e:

            self.logbuf.insert(self.logbuf.get_end_iter(), 'Error:\n' + str(e))

            return

        self.execute_poll()

    def execute_poll(self, source=None, cond=None):

        try:

            state = self.conn.poll()

        except (psycopg2.OperationalError, psycopg2.DatabaseError), e:

            self.logbuf.insert(self.logbuf.get_end_iter(), 'Error:\n' + str(e))

            return

        if state == psycopg2.extensions.POLL_OK:

            self.execute_finish()

        elif state == psycopg2.extensions.POLL_WRITE:

            gobject.io_add_watch(self.conn.fileno(), gobject.IO_OUT, self.execute_poll)

        elif state == psycopg2.extensions.POLL_READ:

            gobject.io_add_watch(self.conn.fileno(), gobject.IO_IN, self.execute_poll)

        else:

            self.logbuf.insert(self.logbuf.get_end_iter(), "poll() returned %s" % state)

            return

    def execute_finish(self):

        t2 = time.time()

        t = (t2 - self.t1)*1000

        self.throbber.clear()

        self.logbuf.insert(self.logbuf.get_end_iter(),

            'Query successful (%d ms, %d rows)\n' % (t, self.curs.rowcount))

        # notices

        for n in self.conn.notices:

            self.logbuf.insert(self.logbuf.get_end_iter(), n)

        if self.curs.rowcount >= 0:

            rows = self.curs.fetchall()

            names = []

            for c in self.curs.description:

                name = c[0]

                typename = self.get_typename(c[1], c[3])

                names += [(name, typename)]

            self.dataview.load_data(names, rows)

        self.tb_execute.set_sensitive(True)

        if self.conn.get_transaction_status() == psycopg2.extensions.TRANSACTION_STATUS_INTRANS:

            self.tb_commit.set_sensitive(True)

            self.tb_rollback.set_sensitive(True)

        else:

            self.tb_connect.set_sensitive(True)

            self.tb_begin.set_sensitive(True)

    def simulate_click(self, tb):

        if tb.get_property('sensitive'):

            tb.get_child().activate()

    def keypress(self, w, event):

        keyname = gtk.gdk.keyval_name(event.keyval)

        if keyname == 'F1':

            self.simulate_click(self.tb_settings)

            return True

        if keyname == 'F2':

            print self.cb_server.popup()

            return True

        if keyname == 'F3':

            print self.cb_dbname.popup()

            return True

        if keyname == 'F4':

            self.simulate_click(self.tb_connect)

            return True

        if keyname == 'F5':

            self.simulate_click(self.tb_execute)

            return True

        if keyname == 'F6':

            self.simulate_click(self.tb_begin)

            return True

        if keyname == 'F7':

            self.simulate_click(self.tb_commit)

            return True

        if keyname == 'F8':

            self.simulate_click(self.tb_rollback)

            return True

        return False

    def toolbar_server_keypress(self, w, event):

        keyname = gtk.gdk.keyval_name(event.keyval)

        if keyname == 'Tab':

            self.populate_db_list()

            self.cb_dbname.grab_focus()

            return True

    def on_server_changed(self, w):

        if self.cb_server.get_active() != -1:

            self.populate_db_list()

    def populate_db_list(self):

        conninfo = self.get_conninfo(True)

        try:

            conn = self.db.get_conn(conninfo)

        except DatabaseError, e:

            self.logbuf.insert(self.logbuf.get_end_iter(), 'Error:\n%s\n' % e)

            return

        curs = conn.cursor()

        curs.execute('SELECT * FROM pg_catalog.pg_database WHERE NOT datistemplate ORDER BY 1')

        psycopg2.extras.wait_select(conn)

        rows = curs.fetchall()

        self.db.put_conn(conninfo, conn)

        for i in range(self.cb_dbname.get_model().iter_n_children(None)):

            self.cb_dbname.remove_text(0)

        for row in rows:

            self.cb_dbname.append_text(row[0])

        self.cb_dbname.set_active(0)

    def reload_server_list(self, sel=None):

        # clean

        for i in range(self.cb_server.get_model().iter_n_children(None)):

            self.cb_server.remove_text(0)

        self.cb_server.get_child().set_text('')

        try:

            # populate

            for server in cfg.servers.server:

                if str(server.name):

                    title = '%s (%s)' % (str(server.name), str(server.host))

                else:

                    title = str(server.host)

                self.cb_server.append_text(title)

        except AttributeError:

            pass

        if not sel is None:

            self.cb_server.set_active(sel)

    def settings(self, w):

        Settings(self)

if __name__ == '__main__':

    try:

        cfg.load('pgconsole.xml.gz')

    except IOError:

        cfg.new('pgconsole.xml.gz')

    try:

        gtkrc = cfg.root.style.gtkrc.text

    except AttributeError:

        pass

    else:

        gtk.rc_parse_string(gtkrc)

    app = PgConsoleApp()

    app.main()