"""Elasticsearch backfill script.

To run type:
DD_MYSQL_HOST=<host> DD_MYSQL_USER=<db_user> DD_MYSQL_PASSWORD=<db_password>
python es_backfill.py --step=100 --id-from=1 --id-to=2 --wait=10 --log-str='.'

"""

import argparse
import functools
import logging
import time

import tenacity

from vectororder.connectors.mysql import dd_db_connector

MYSQL_INT_SIGNED_MAX_VALUE = 2147483647
RETRY_ATTEMPTS = 100
RETRY_MAX_WAIT_TIME = 120

logger = logging.getLogger(__name__)

retry = functools.partial(
    tenacity.retry,
    stop=tenacity.stop_after_attempt(RETRY_ATTEMPTS),
    wait=tenacity.wait_exponential(multiplier=2, max=RETRY_MAX_WAIT_TIME),
    before_sleep=tenacity.before_sleep_log(logger, logging.WARNING),
    reraise=True,
)


@retry()
def get_next_id(start_id, step):
    """Get ID to update to.

    Args:
        start_id (int): ID to update from
        step (int): update records count

    Returns:
        tuple: maximum int ID to update and a flag indicating this was the end
        of the table.
    """
    with dd_db_connector.db_session() as session:
        # The UNION is needed because with some offset and limit combinations
        # we could miss the latest item(s).
        result = session.execute(
            """
            SELECT eqd.encoding_queue_detail_id AS eqd_id
            FROM encoding_queue_detail eqd
            WHERE eqd.encoding_queue_detail_id >= :start_id
            LIMIT 1 OFFSET :offset
            UNION
            SELECT MAX(max_eqd.encoding_queue_detail_id) AS eqd_id
            FROM encoding_queue_detail max_eqd;
            """, {'start_id': start_id, 'offset': step - 1}
        )
        # There are always one or two rows in result:
        #  - largest ID from the batch (if found)
        #  - the largest ID in the table that we want to update
        #    (is always there)
        # We need the first one as the next ID. The absence of the second one
        # indicates the fact that we've reached the end of the table.
        all_rows = list(result.fetchall())
        is_highest_id = len(all_rows) == 1
        return all_rows[0][0], is_highest_id


@retry()
def update_records(start_id, end_id, log_str):
    """Update a set of rows.

    Args:
        start_id (int): ID to update from in the current batch
        end_id (int): ID to update to in the current batch
        log_str (str): str to add to log
    """
    print('Updating IDs {} - {}'.format(start_id, end_id), flush=True)
    with dd_db_connector.db_session() as session:
        result = session.execute(
            """
                UPDATE encoding_queue_detail
                SET
                    error_log = CONCAT(IFNULL(error_log, ""), :log_str),
                    last_updated = last_updated
                WHERE encoding_queue_detail_id >= :start_id
                    AND encoding_queue_detail_id <= :end_id
            """, {'log_str': log_str, 'start_id': start_id, 'end_id': end_id})
        print(
            f'Records updated in this iteration: {result.rowcount:d}',
            flush=True)
        return result.rowcount


def execute_batch_update(*, id_from, id_to, step, wait, log_str):
    """Execute batch update loop.

    Args:
        id_from (int): ID to update from
        id_to (int): ID to update to
        step (int): records number to update in a single SQL request
        wait (int): sleep interval (s) after each update
        log_str (str): str to add to log
    """
    if id_to <= id_from:
        raise ValueError('id_from should be less than id_to')

    start_id = id_from
    end_of_table = False
    maximum_id_reached = False
    records_updated = 0
    while not end_of_table and not maximum_id_reached:
        end_id, end_of_table = get_next_id(start_id, step)
        if end_id >= id_to:
            # We're approaching the maximum ID we were supposed to update.
            # This is going to be the last iteration.
            print(f'Limiting end_id to {id_to}.', flush=True)
            end_id = id_to
            maximum_id_reached = True
        records_updated += update_records(start_id, end_id, log_str)
        print(f'Total records updated: {records_updated:d}', flush=True)
        start_id = end_id + 1
        time.sleep(wait)


def setup_parser():
    """Set args parser."""
    new_parser = argparse.ArgumentParser()
    new_parser.add_argument(
        '--id-from', required=False, default=0, type=int, dest='id_from',
        help='First encoding_queue_id to update')
    new_parser.add_argument(
        '--id-to', required=False, default=MYSQL_INT_SIGNED_MAX_VALUE,
        type=int, dest='id_to',
        help='Last encoding_queue_id to update')
    new_parser.add_argument(
        '--step', required=False, default=10000, type=int,
        help='Number of records to update per one query')
    new_parser.add_argument(
        '--wait', required=False, default=0, type=int,
        help='Number of seconds to wait after each update')
    new_parser.add_argument(
        '--log-str', required=False, default=' ', type=str, dest='log_str',
        help='String to concat to error_log to force update')
    return new_parser


if __name__ == '__main__':
    parser = setup_parser()
    args = parser.parse_args()
    execute_batch_update(
        id_from=args.id_from,
        id_to=args.id_to,
        step=args.step,
        wait=args.wait,
        log_str=args.log_str,
    )