#!/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()