from pgtoolkit import pgmanager, pgbrowser

from pycolib.configparser import ConfigParser
from pycolib.coloredformatter import ColoredFormatter
from pycolib.ansicolor import highlight

import argparse
import logging
import re
import textwrap


class ConnectionInfoNotFound(Exception):
    pass


class BadArgsError(Exception):
    pass


class ToolDescriptionFormatter(argparse.HelpFormatter):
    """Help message formatter which retains any formatting in descriptions."""

    def _fill_text(self, text, width, indent):
        return textwrap.dedent(text)


class ToolBase:

    def __init__(self, name, desc=None, **kwargs):
        self.config = ConfigParser()
        self.parser = argparse.ArgumentParser(prog=name, description=desc or self.__doc__,
            formatter_class=ToolDescriptionFormatter)
        self.pgm = pgmanager.get_instance()
        self.target_isolation_level = None

    def setup(self, args=None):
        self.specify_args()
        self.load_args(args)
        self.init_logging()

    def specify_args(self):
        self.config.add_option('databases', dict)
        self.config.add_option('meta_db')
        self.config.add_option('meta_query')
        self.parser.add_argument('-Q', dest='show_queries', action='store_true',
            help='Print database queries.')
        self.parser.add_argument('-C', dest='config_file', type=str,
            help='Additional config file (besides pgtoolkit.conf).')

    def load_args(self, args=None, config_file=None):
        # Parse command line arguments
        self.args = self.parser.parse_args(args)
        # Load global config
        self.config.load('/etc/pgtoolkit.conf', must_exist=False)
        # Load local config
        self.config.load(config_file or 'pgtoolkit.conf', must_exist=False)
        # Load additional config
        if self.args.config_file:
            self.config.load(self.args.config_file)

    def init_logging(self):
        # logging
        format = ColoredFormatter(highlight(1,7,0)+'%(asctime)s %(levelname)-5s'+highlight(0)+' %(message)s', '%H:%M:%S')
        handler = logging.StreamHandler()
        handler.setFormatter(format)
        handler.setLevel(logging.DEBUG)
        self.log = logging.getLogger('main')
        self.log.addHandler(handler)
        self.log.setLevel(logging.DEBUG)

        log_notices = logging.getLogger('pgmanager_notices')
        log_notices.addHandler(handler)
        log_notices.setLevel(logging.DEBUG)

        if self.args.show_queries:
            log_sql = logging.getLogger('pgmanager_sql')
            log_sql.addHandler(handler)
            log_sql.setLevel(logging.DEBUG)

    def prepare_conn_from_metadb(self, name, lookup_name):
        """Create connection in pgmanager using meta DB.

        name -- Name for connection in pgmanager.
        lookup_name -- Name of connection in meta DB.

        """
        if not self.pgm.knows_conn('meta'):
            self.pgm.create_conn(name='meta', dsn=self.config.meta_db)
        with self.pgm.cursor('meta') as curs:
            curs.execute(self.config.meta_query, [lookup_name])
            row = curs.fetchone_dict()
            curs.connection.commit()
            if row:
                self.pgm.create_conn(name=name,
                    isolation_level=self.target_isolation_level,
                    **row)
                return True
        self.pgm.close_conn('meta')

    def prepare_conn_from_config(self, name, lookup_name):
        """Create connection in pgmanager using info in config.databases."""
        if self.config.databases:
            if lookup_name in self.config.databases:
                dsn = self.config.databases[lookup_name]
                self.pgm.create_conn(name=name,
                    isolation_level=self.target_isolation_level,
                    dsn=dsn)
                return True

    def prepare_conns(self, **kwargs):
        """Create connections in PgManager.

        Keyword arguments meaning:
            key: connection name for use in PgManager
            value: connection name in config or meta DB

        """
        for name in kwargs:
            lookup_name = kwargs[name]
            found = self.prepare_conn_from_config(name, lookup_name)
            if not found and self.config.meta_db:
                found = self.prepare_conn_from_metadb(name, lookup_name)
            if not found:
                raise ConnectionInfoNotFound('Connection name "%s" not found in config nor in meta DB.' % lookup_name)


class SimpleTool(ToolBase):

    def __init__(self, name, desc=None, **kwargs):
        ToolBase.__init__(self, name, desc, **kwargs)

    def specify_args(self):
        ToolBase.specify_args(self)
        self.config.add_option('target', type=str, default=None)
        self.parser.add_argument('target', nargs='?', type=str, help='Target database')

    def load_args(self, args=None, config_file=None):
        ToolBase.load_args(self, args, config_file)
        self.target = self.args.target or self.config.target or 'default'

    def setup(self, args=None):
        ToolBase.setup(self, args)
        self.prepare_conns(target=self.target)


class SrcDstTool(ToolBase):

    def __init__(self, name, desc=None, *, allow_reverse=False, force_reverse=False, **kwargs):
        ToolBase.__init__(self, name, desc, **kwargs)
        self.allow_reverse = allow_reverse
        self.force_reverse = force_reverse

    def specify_args(self):
        ToolBase.specify_args(self)
        self.parser.add_argument('src', metavar='source', type=str, help='Source database')
        self.parser.add_argument('dst', metavar='destination', type=str, help='Destination database')
        if self.allow_reverse:
            self.parser.add_argument('-r', '--reverse', action='store_true', help='Reverse operation. Swap source and destination.')

    def load_args(self, args=None, config_file=None):
        ToolBase.load_args(self, args, config_file)
        if self.is_reversed():
            self.args.src, self.args.dst = self.args.dst, self.args.src

    def setup(self, args=None):
        ToolBase.setup(self, args)
        self.prepare_conns(src=self.args.src, dst=self.args.dst)

    def is_reversed(self):
        return ('reverse' in self.args and self.args.reverse) or self.force_reverse


class SrcDstTablesTool(SrcDstTool):

    def specify_args(self):
        SrcDstTool.specify_args(self)
        self.parser.add_argument('-t', '--src-table', metavar='source_table',
            dest='srctable', type=str, default='', help='Source table name.')
        self.parser.add_argument('-s', '--src-schema', metavar='source_schema',
            dest='srcschema', type=str, default='', help='Source schema name (default=public).')
        self.parser.add_argument('--dst-table', metavar='destination_table',
            dest='dsttable', type=str, default='', help='Destination table name (default=source_table).')
        self.parser.add_argument('--dst-schema', metavar='destination_schema',
            dest='dstschema', type=str, default='', help='Destination schema name (default=source_schema).')
        self.parser.add_argument('--regex', action='store_true', help="Use RE in schema or table name.")

    def load_args(self, args=None, config_file=None):
        SrcDstTool.load_args(self, args, config_file)
        self.load_table_names()

    def load_table_names(self):
        self.schema1 = self.args.srcschema
        self.table1 = self.args.srctable
        self.schema2 = self.args.dstschema
        self.table2 = self.args.dsttable

        # check regex - it applies to source name, dest name must not be specified
        # applies to only one - schema or table name
        if self.args.regex:
            if self.table2 or (self.schema2 and not self.table1):
                raise BadArgsError('Cannot specify both --regex and --dst-schema, --dst-table.')
        # schema defaults to public
        if self.table1 and not self.schema1:
            self.schema1 = 'public'
        # dest defaults to source
        if not self.schema2:
            self.schema2 = self.schema1
        if not self.table2:
            self.table2 = self.table1

        # swap src, dst when in reverse mode
        if self.is_reversed():
            self.schema1, self.schema2 = self.schema2, self.schema1
            self.table1, self.table2 = self.table2, self.table1

    def tables(self):
        '''Generator. Yields schema1, table1, schema2, table2.'''
        srcconn = self.pgm.get_conn('src')
        try:
            srcbrowser = pgbrowser.PgBrowser(srcconn)
            if self.args.regex:
                if not self.table1:
                    # all tables from schemas defined by regex
                    for item in self._iter_schemas_regex(srcbrowser, self.schema1):
                        yield item
                else:
                    # all tables defined by regex
                    for item in self._iter_tables_regex(srcbrowser, self.schema1, self.schema2, self.table1):
                        yield item
            else:
                if not self.table1:
                    if not self.schema1:
                        # all tables from all schemas
                        for item in self._iter_schemas_regex(srcbrowser, self.schema1):
                            yield item
                    else:
                        # all tables from specified schema
                        for item in self._iter_tables_regex(srcbrowser, self.schema1, self.schema2, self.table1):
                            yield item
                else:
                    # one table
                    yield (self.schema1, self.table1, self.schema2, self.table2)
        finally:
            self.pgm.put_conn(srcconn, 'src')

    def _iter_schemas_regex(self, browser, regex):
        for schema in browser.list_schemas():
            if schema['system']:
                continue
            schemaname = schema['name']
            if re.match(regex, schemaname):
                for item in self._iter_tables_regex(browser, schemaname, schemaname, ''):
                    yield item

    def _iter_tables_regex(self, browser, schema1, schema2, regex):
        for table in browser.list_tables(schema1):
            tablename = table['name']
            if re.match(regex, tablename):
                yield (schema1, tablename, schema2, tablename)