#!/usr/bin/env python3

# -*- coding: utf-8 -*-

import locale

locale.setlocale(locale.LC_ALL, '')

from tuikit import *

from pgtoolkit import pgbrowser

from pgtoolkit.toolbase import SimpleTool

class MyApplication(Application, SimpleTool):

    PLACEHOLDER = '__placeholder__'

    def __init__(self):

        Application.__init__(self)

        SimpleTool.__init__(self, name='browser', desc='PostgreSQL database browser.')

        self.init()

        self.top.connect('keypress', self.on_top_keypress)

        self.browser = pgbrowser.PgBrowser(self.pgm.get_conn('target'))

        self.model = TreeModel()

        self.view = TreeView(self.model)

        self.view.connect('expand', self.on_view_expand)

        self.top.add(self.view)

        vert = VerticalLayout()

        self.top.layout(vert)

        self.populate_schemas()

    def populate_schemas(self):

        schemas = self.browser.list_schemas()

        schema_list = [schema['name'] for schema in schemas if not schema['system']]

        self.model.add('/', TreeNode('schemas', 'Schemas (%d)' % len(schema_list)))

        self.model.add('/schemas', schema_list)

        for schema_name in schema_list:

            path = '/schemas/' + schema_name

            self.model.add(path, self.PLACEHOLDER)

            self.view.collapse(path)

    def populate_tables(self, node, schema_name):

        """Read table list, put in node.

        Node is replaced by 'Tables (N)' node,

        tables are created as its children.

        """

        tables = self.browser.list_tables(schema=schema_name)

        table_list = [table['name'] for table in tables]

        node.name = 'tables'

        node.title = 'Tables (%d)' % len(table_list)

        for table_name in table_list:

            node.add(TreeNode(table_name))

    def on_view_expand(self, node):

        if node.parent.name == 'schemas' and node.children[0].name == self.PLACEHOLDER:

            schema_name = node.name

            self.populate_tables(node.children[0], schema_name)

    def on_top_keypress(self, keyname, char):

        if keyname == 'escape':

            self.terminate()

if __name__ == '__main__':

    app = MyApplication()

    app.start()