"""Kick-off script for bulk-assets-ingester orchestration."""

from connectors.logger import logger as log
import constants
from math import ceil
from time import sleep

import shortuuid
import uuid

from utils import snowflake_utils as snow_util

from sql.snowflake.copy_into_s3_as_json import COPY_INTO_S3_AS_JSON
from connectors.sentry import sentry_wrap
from connectors.sfn import get_full_execution_count_by_arn
from connectors.sfn import get_full_execution_name_list_by_arn
from constants import SFN_START_DELAY
import config

FILE_EXTENSION = 'json'


def write_output_files(file_name, per_file_rows, asset_count):
    """Write out the json files."""
    offset = 0
    while offset < asset_count:
        start = str(offset).zfill(4)
        end = str(min(asset_count, offset + per_file_rows - 1)).zfill(4)
        loop_file_name = f'{file_name}_{start}-{end}.{FILE_EXTENSION}'

        bucket = config.OUTPUT_S3_BUCKET
        location = config.OUTPUT_S3_LOCATION

        # Log out the name of the file created.
        log.info(f'Writing s3://{bucket}/{location}/{loop_file_name}')

        # Write out <machine_count> files using offsets
        snow_util.get_snowflake_results(
            sql=COPY_INTO_S3_AS_JSON,
            params={
                's3_bucket': bucket,
                's3_location': location,
                'file_name': loop_file_name,
                'source_snowflake_database': config.SNOWFLAKE_DATABASE,
                'source_snowflake_schema': config.SNOWFLAKE_SCHEMA,
                'snowflake_table_name': config.SOURCE_TABLE_NAME,
                'limit': per_file_rows,
                'offset': offset,
                'aws_key_id': config.AWS_ACCESS_KEY_ID,
                'aws_secret_id': config.AWS_SECRET_ACCESS_KEY
            }
        )

        offset = offset + per_file_rows

        log.info(f'Sleeping for {SFN_START_DELAY} secs.')
        sleep(SFN_START_DELAY)


@sentry_wrap
def main():
    """Main script execution function."""
    log.info('Beginning Bulk Asset Uploader.')

    # Create key
    fingerprint = shortuuid.encode(uuid.uuid4())

    # Create file name
    file_name = \
        f'{config.SOURCE_TABLE_NAME}_{fingerprint[:7]}'

    # Get the count of current executions
    current_sfn_count = get_full_execution_count_by_arn(
        config.BULK_ASSETS_INGESTER_SFN_ARN, 'RUNNING'
    )

    # Report running sfn count
    msg = f'{current_sfn_count} currently running instances of ' \
          f'{config.BULK_ASSETS_INGESTER_SFN_NAME}.'
    log.info(msg)

    current_exec_name_list = get_full_execution_name_list_by_arn(
        config.BULK_ASSETS_INGESTER_SFN_ARN, 'RUNNING'
    )

    # Get the count of rows in the input table
    asset_count = snow_util.get_count_all_rows(table=config.SOURCE_TABLE_NAME)

    # Report asset count
    msg = f'{asset_count} assets found in table {config.SOURCE_TABLE_NAME}'
    log.info(msg)

    # Save number of state machines to invoke
    sfn_count = config.CONCURRENT_STATE_MACHINE_COUNT

    # Calculate per-file rows
    per_file_rows = ceil(asset_count / int(sfn_count))

    # Report per-file rows
    msg = f'{per_file_rows} will be written to each table'
    log.info(msg)

    # Auto shape state machine count
    i = int(config.CONCURRENT_STATE_MACHINE_COUNT)
    while per_file_rows < constants.PER_FILE_ROW_MIN and i > 1:
        log.info(
            f'The number of rows delegated to each state '
            f'machine does not meet per-file minimum of '
            f'{constants.PER_FILE_ROW_MIN}. '
            f'decreasing CONCURRENT_STATE_MACHINE_COUNT of '
            f'{str(i)} to {str(i-1)}.')
        i = i - 1
        per_file_rows = ceil(asset_count / i)

    new_machine_count = i

    # Check if you can run more State machines (too many concurrencies)
    count_ceil = current_sfn_count + new_machine_count
    if count_ceil > constants.MAX_STATE_MACHINES:
        raise RuntimeError(f'The number of concurrent state machines '
                           f'requested ({new_machine_count}) exceeds the '
                           f'allowed maximum of '
                           f'{constants.MAX_STATE_MACHINES}. Please decrease '
                           f'the CONCURRENT_STATE_MACHINE_COUNT below '
                           f'{constants.MAX_STATE_MACHINES-current_sfn_count} '
                           f'to proceed.')

    # Check if too many rows in each file
    if per_file_rows > constants.PER_FILE_ROW_LIMIT:
        raise RuntimeError(f'The number of rows delegated to each state '
                           f'machine exceeds {constants.PER_FILE_ROW_LIMIT}.'
                           f'Please increase the '
                           f'CONCURRENT_STATE_MACHINE_COUNT or decrease the '
                           f'number of rows in the input table to proceed.')

    # Write out S3 files.
    log.info(f'Launching {new_machine_count} state machines.')
    # DEBUG: Disable file output for testing
    # write_output_files(file_name, per_file_rows, asset_count)

    # TODO: Can we add tuning params within this function

    # Sleeping to let sfn start.
    log.info(f'Sleeping for {SFN_START_DELAY} secs.')
    sleep(SFN_START_DELAY)

    # Get the list of current executions
    full_exec_name_list = get_full_execution_name_list_by_arn(
        config.BULK_ASSETS_INGESTER_SFN_ARN, 'RUNNING'
    )

    # Report running sfn count
    msg = f'{len(full_exec_name_list)} currently running instances of ' \
          f'{config.BULK_ASSETS_INGESTER_SFN_NAME}.'
    log.info(msg)

    new_exec_name_list = [
        x for x in full_exec_name_list if x not in current_exec_name_list]

    for name in new_exec_name_list:
        log.info(f'Execution {name} was started. '
                 f'https://console.aws.amazon.com/states/home?region=us-east'
                 f'-1#/executions/details/arn:aws:states:us-east-1:437795906'
                 f'767:execution:qa-bulk-assets-ingester-sfn:{name}')

    # Report new sfn count
    msg = f'{len(new_exec_name_list)} new instances of ' \
          f'{config.BULK_ASSETS_INGESTER_SFN_NAME} started.'
    log.info(msg)


if __name__ == '__main__':
    main()