from pgtoolkit.toolbase import SimpleTool

import logging.handlers

import time

from datetime import datetime, timedelta

class LoopQueryTool(SimpleTool):

    """

    Execute queries in loop, with configurable interval.

    """

    def __init__(self):

        SimpleTool.__init__(self, name='loopquery')

        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('--mins', dest='delay_mins', type=int, help='Delay between queries in minutes.')

        self.parser.add_argument('--secs', dest='delay_secs', type=int, help='Delay between queries in seconds.')

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

        self.config.add_option('delay_mins', type=int, default=0)

        self.config.add_option('delay_secs', type=int, default=0)

        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

        self.delay_mins = self.args.delay_mins or self.config.delay_mins

        self.delay_secs = self.args.delay_secs or self.config.delay_secs

    def init_logging(self):

        SimpleTool.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):

        self.reset()

        while True:

            self.wait()

            self.action()

    def reset(self):

        """Check current time, set next action time."""

        dt = datetime.today()

        dt = dt.replace(microsecond = 0)

        self.next_action_time = dt + timedelta(minutes = self.delay_mins,

                                               seconds = self.delay_secs)

    def wait(self):

        """Wait for action time, compute next action time."""

        now = datetime.today()

        self.log.debug('Next run %s', self.next_action_time)

        if self.next_action_time > now:

            td = self.next_action_time - now

            self.log.debug('Sleep %ds', td.seconds + td.microseconds/1e6)

            time.sleep(td.seconds + td.microseconds/1e6)

        self.next_action_time += timedelta(minutes = self.delay_mins,

                                           seconds = self.delay_secs)

        # in case that action took too long and next planned time would

        # be in past -> reset planner

        if self.next_action_time < now:

            self.reset()

    def action(self):

        """Execute the queries."""

        for q in self.queries:

            self.log.info('%s', q)

            with self.pgm.cursor('target') as curs:

                curs.execute(q)

        self.log.info('Done')

cls = LoopQueryTool