from pgtoolkit.toolbase import SimpleTool

import logging.handlers
import time
from datetime import datetime, timedelta
from psycopg2 import ProgrammingError


class RunQueryTool(ToolBase):

    """
    Execute configured queries in target database.
    """

    def __init__(self):
        SimpleTool.__init__(self, name='runquery')
        self.target_isolation_level = 'autocommit'

    def specify_args(self):
        SimpleTool.specify_args(self)
        self.parser.add_argument('-q', dest='queries', metavar='QUERY', nargs='*', help='Queries to run.')
        self.parser.add_argument('-f', dest='file', metavar='FILE', help='Read query from file.')
        self.parser.add_argument('--one-query-per-line', action='store_true', help='When reading queries from file, consider each line as separate query.')
        self.parser.add_argument('-p', '--parameter', dest='parameters', metavar='PARAM=VALUE', nargs='*',
            help="If query should be used as format template, these parameters will be substituted.")
        self.parser.add_argument('--output-file', dest='output_file', metavar='OUTPUT_FILE', help='Write query result in specified file.')
        self.parser.add_argument('--format', dest='format', metavar='FORMAT', help='Format string for each line in output file (using Python\'s format()).')

        self.config.add_option('queries', type=list, default=[])
        self.config.add_option('log_path', type=str)

    def load_args(self, args=None, config_file=None):
        SimpleTool.load_args(self, args, config_file)
        self.queries = self.args.queries or self.config.queries
        # read query from file
        if self.args.file:
            with open(self.args.file, 'r', encoding='utf8') as f:
                data = f.read()
            if self.args.one_query_per_line:
                file_queries = [ln for ln in data.splitlines() if not ln.lstrip().startswith('--')]
                self.queries = file_queries + self.queries
            else:
                self.queries.insert(0, data)
        # prepare parameters
        self._prepare_parameters(self.args.parameters)

    def init_logging(self):
        ToolBase.init_logging(self)
        if self.config.log_path:
            self.init_file_logs(self.config.log_path)

    def init_file_logs(self, path):
        format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S')
        handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5)
        handler.setFormatter(format)
        handler.setLevel(logging.DEBUG)
        logging.getLogger('main').addHandler(handler)

        format = logging.Formatter('%(asctime)s %(message)s', '%y-%m-%d %H:%M:%S')
        handler = logging.handlers.TimedRotatingFileHandler(path+'/pgnotices.log', when='midnight', backupCount=5)
        handler.setFormatter(format)
        handler.setLevel(logging.DEBUG)
        logging.getLogger('pgmanager_notices').addHandler(handler)

    def main(self):
        """Execute the queries."""
        for q in self.queries:
            if self.parameters:
                q = q.format(**self.parameters)
            self.log.info('%s', q if len(q) < 100 else q[:100]+'...')
            with self.pgm.cursor('target') as curs:
                curs.execute(q)
                self.log.info('Rows affected: %d', curs.rowcount)
                try:
                    rows = curs.fetchall_dict()
                    self._write_output_file(rows)
                except ProgrammingError:
                    pass
        self.log.info('Done')

    def _write_output_file(self, rows):
        if not self.args.output_file:
            return
        with open(self.args.output_file, 'w', encoding='utf8') as f:
            for row in rows:
                print(self.args.format.format(row), file=f)

    def _prepare_parameters(self, parameters):
        self.parameters = {}
        for parameter in parameters or ():
            name, value = parameter.split('=', 1)
            self.parameters[name] = value


cls = RunQueryTool