#!/usr/bin/env python3
#
# Copy data between tables with same table schema.
#
# Copies full table, target table must be empty.
# Can copy multiple tables in one run.
# Sorts the tables according to references.
#
# This may be used instead of dump/restore.
#

import io

from pgtoolkit import toolbase, pgmanager, pgbrowser, pgdatacopy
from pgtoolkit.progresswrapper import ProgressWrapper


class TableCopyTool(toolbase.SrcDstTablesTool):
    def __init__(self):
        toolbase.SrcDstTablesTool.__init__(self, name='tablecopy', desc='Table copy tool.')

        self.parser.add_argument('-n', '--no-action', dest='noaction', action='store_true',
            help="Do nothing, just print tables to be copied. Useful in combination with --regex.")
        self.parser.add_argument('--no-sort', dest='nosort', action='store_true',
            help="Do not sort. By default, tables are sorted by foreign key references.")
        self.parser.add_argument('--disable-triggers', dest='notriggers', action='store_true',
            help="Disable all triggers for the session.")

        self.init()

    def main(self):
        self.srcconn = self.pgm.get_conn('src')
        self.dstconn = self.pgm.get_conn('dst')

        if self.args.notriggers:
            curs = self.dstconn.cursor()
            curs.execute('SET session_replication_role = replica;')
            curs.close()
            self.dstconn.commit()

        dc = pgdatacopy.PgDataCopy(self.srcconn, self.dstconn)

        if self.args.nosort:
            for table in self.tables():
                self.copy_table(dc, *table)
        else:
            # sort tables with respect to references
            details = dict()
            pending = set()
            references = dict()

            # build list of all table to be copied (pending) and references map
            for table in self.tables():
                srcschema, srctable, dstschema, dsttable = table
                name = dstschema + '.' + dsttable
                details[name] = table
                pending.add(name)
                references[name] = self.get_references(dstschema, dsttable)

            # copy files with fulfilled references, repeat until all done
            while pending:
                for name in list(pending):
                    ok = True
                    for ref in references[name]:
                        if ref in pending:
                            ok = False
                    if ok:
                        table = details[name]
                        self.copy_table(dc, *table)
                        pending.remove(name)

    def get_references(self, schema, table):
        browser = pgbrowser.PgBrowser(self.dstconn)
        cons = browser.list_constraints(table, schema)
        return [con['fschema'] + '.' + con['fname'] for con in cons if con['type'] == 'f']

    def copy_table(self, dc, srcschema, srctable, dstschema, dsttable):
        print('Copying [%s] %s.%s --> [%s] %s.%s' % (
            self.args.src, srcschema, srctable,
            self.args.dst, dstschema, dsttable))

        if self.args.noaction:
            return

        dc.set_source(srctable, srcschema)
        dc.set_destination(dsttable, dstschema)

        try:
            dc.check()
        except pgdatacopy.TargetNotEmptyError as e:
            print(' - error:', str(e))
            return

        print(' - read                           ')
        buf = io.BytesIO()
        wrapped = ProgressWrapper(buf)
        dc.read(wrapped)
        data = buf.getvalue()
        buf.close()

        print(' - write                          ')
        buf = io.BytesIO(data)
        wrapped = ProgressWrapper(buf, len(data))
        dc.write(wrapped)
        buf.close()

        print(' - analyze                        ')
        dc.analyze()


tool = TableCopyTool()
tool.main()