"""Dimension Refresh Flow.

========================
A Garcon workflow for updating Snowflake dimension tables.
A typical dimension refresh involves:
1) Bootstrapping dimension specific variables.
2) Exporting the dimension source data into s3. The source is typically from
    MySQL tables in the art_relations db, but can also be from other Snowflake
    dimension tables.
3) Loading a staging table in Snowflake.
4) Updating existing dimensions in Snowflake.
5) Inserting new dimensions into Snowflake.
6) Updating any related dimensions tables.

Dimensions are defined in dim_refresh_etl/conf/dim/<dimension>/yml.

Params:
    dimension (str): The dimension.
"""

import logging
import os

import boto3
from botocore.config import Config
from garcon import activity
from garcon import runner
from garcon.param import StaticParam
from garcon_contrib.aws import garcon_aws_cli_tool
from garcon_contrib.aws import garcon_s3
from garcon_contrib.aws import garcon_sns
from garcon_contrib.gzip import garcon_gzip
from garcon_contrib.mysql import garcon_mysql
from garcon_contrib.pipe import garcon_pipe
from garcon_contrib.pipe import garcon_pipe_task_runner

from dim_refresh_etl import conf
from dim_refresh_etl.conf.config import SF_CONFIG
from dim_refresh_etl.tasks import bootstrap
from dim_refresh_etl.tasks import db
from dim_refresh_etl.tasks import downstream
from dim_refresh_etl.tasks import notification
from dim_refresh_etl.util import environment, sentry_utils

logger = logging.getLogger('dim_refresh_etl')

ART_RELATIONS_DB = conf.get_static_dict('mysql', 'art_relations')
SNS_CONFIG = conf.get_static_dict('aws', 'sns')
SWF_PROD_DOMAIN = 'prod_dim_refresh'
SWF_QA_DOMAIN = 'qa_dim_refresh'


class Flow:
    """Refresh flow class."""

    timeout = 60 * 60 * 6  # 6 hours

    def __init__(self):
        """Create a Dimension WorkFlow flow."""
        if environment.name == environment.PROD:
            self.domain = SWF_PROD_DOMAIN
        elif environment.name == environment.QA:
            self.domain = SWF_QA_DOMAIN
        else:
            self.domain = os.getenv('DEV_DOMAIN', 'dev')

        self.name = 'dim_refresh'
        self.version = '1.1'
        self.client = boto3.client('swf', config=Config(
            connect_timeout=60,
            region_name='us-east-1',
            read_timeout=180,
            retries={'max_attempts': 2}
        ))
        self.create = activity.create(
            self.client, self.domain, self.name, version=self.version,
            on_exception=self.on_exception)

    @property
    def conf_aws(self):
        """Get AWS credentials.

        Returns:
            dict: aws credentials.
        """
        return conf.getconf('aws')['aws']

    @property
    def conf_snowflake(self):
        """Get Snowflake credentials.

        Returns:
            dict: Snowflake credentials.
        """
        return SF_CONFIG

    @property
    def conf_dim_snowflake_options(self):
        """Get Snowflake options for COPY INTO.

        For dimension related commands.

        Returns:
            dict: contains FILE_FORMAT string for COPY statement.
        """
        return conf.getconf('snowflake')['dim']

    def decider(self, schedule, context):
        """Flow decider.

        Arg:
            schedule (callable): Call the scheduler.
        """
        bootstrap = schedule('bootstrap', self.bootstrap_activity)

        # if the bootstrap hasn't loaded, we don't know our source db
        if not bootstrap.ready:
            logger.debug('Not ready to decide on source db')
            return

        clean_s3_source_export = schedule(
            'clean_s3_source_export',
            self.clean_s3_source_export_activity,
            requires=[bootstrap])

        # check what sort of source task needs to get scheduled
        if bootstrap.result.get('dim_refresh.source_db') == 'art_relations':
            export_source_table = schedule(
                'export_mysql_source_table',
                self.export_mysql_source_activity,
                requires=[clean_s3_source_export])
        elif bootstrap.result.get('dim_refresh.source_db') == 'snowflake':
            export_source_table = schedule(
                'export_snowflake_source_table',
                self.export_snowflake_source_activity,
                requires=[clean_s3_source_export])
        else:
            raise ValueError('Unrecognized source_db {}'.format(
                bootstrap.result.get('dim_entity.source_db')))

        load_snowflake_staging_table = schedule(
            'load_snowflake_staging_table',
            self.load_snowflake_staging_table_activity,
            requires=[export_source_table])

        update_existing_dimensions = schedule(
            'update_existing_dimensions',
            self.update_existing_dimensions_activity,
            requires=[load_snowflake_staging_table])

        insert_new_dimensions = schedule(
            'insert_new_dimensions',
            self.insert_new_dimensions_activity,
            requires=[update_existing_dimensions])

        update_new_dimensions = schedule(
            'update_new_dimensions',
            self.update_new_dimensions_activity,
            requires=[insert_new_dimensions])

        delete_old_dimensions = schedule(
            'delete_old_dimensions',
            self.delete_old_dimensions_activity,
            requires=[update_new_dimensions])

        related_dimension_updates = schedule(
            'related_dimension_updates',
            self.related_dimension_updates_activity,
            requires=[delete_old_dimensions])

        notify_dimension_refresh = schedule(
            'notify_dimension_refresh',
            self.notify_dimension_refresh_activity,
            requires=[related_dimension_updates])

        if 'downstream_contexts' in context:
            schedule(
                'downstream_dim_flows',
                self.downstream_dim_flows_activity,
                requires=[notify_dimension_refresh])

    def on_exception(self, actor, exception):
        """Capture an exception that has occurred in the application.

        Args:
            actor (Activity, DeciderWorker): the actor that has received
                the exception.
            exception (Exception): the exception to capture.
        """
        # client grabs sentry dns from SENTRY_DSN environment variable
        sentry_utils.capture_exception()

        if isinstance(actor, activity.Activity):
            actor.logger.error(exception, exc_info=True)
        else:
            logger.error(exception, exc_info=True)

    def workflow_id(self, context):
        """Generate the workflow id.

        Args:
            context (dict): initial context of the workflow. Must have
                the 'dim_type' key set.
        """
        assert context.get('dim_type'), "Initial context must set 'dim_type'"
        workflow_id = '{}'.format(context.get('dim_type'))
        return workflow_id

    @property
    def bootstrap_activity(self):
        """Bootstrap the configuration.

        Get the initial values from the context and hydrate context values used
        in the rest of the flow.
        """
        return self.create(
            name='bootstrap',
            retry=10,
            tasks=runner.Sync(
                bootstrap.bootstrap_task.fill(dim_type='dim_type')))

    @property
    def clean_s3_source_export_activity(self):
        """Cleanup any s3 data if dimension update had run previously.

        This is the exception to the norm, & might cause funky results bc
        of s3 eventual consistency. If multiple runs becomes a regular thing,
        we need to namespace each run with a md5 hash or timestamp.
        """
        return self.create(
            name='clean_s3_source_export',
            tasks=runner.Sync(
                garcon_s3.remove_files_from_path.fill(
                    path='dim_refresh.staging_path')))

    @property
    def export_mysql_source_activity(self):
        """Export dimension source from MySQL.

        Export dimension source input from MySQL to a gzip file in s3.
        """
        return self.create(
            name='export_mysql_source_table',
            tasks=garcon_pipe_task_runner.Pipe(
                db.art_relation_sql_to_stdout.fill(
                    namespace='pipe_sql_to_stdout',
                    query='dim_refresh.export_sql'),
                garcon_mysql.pipe_mysql_from_stdin_to_stdout.fill(
                    namespace='mysql',
                    pipe='pipe_sql_to_stdout.pipe',
                    **ART_RELATIONS_DB),
                garcon_gzip.pipe_gzip_from_stdin_to_stdout.fill(
                    namespace='gzip',
                    pipe='mysql.pipe'),
                garcon_aws_cli_tool.upload_to_s3_from_stdin.fill(
                    namespace='s3',
                    pipe='gzip.pipe',
                    destination_s3_key='dim_refresh.staging_key',
                    bucket='dim_refresh.staging_bucket'),
                garcon_pipe.communicate.fill(
                    pipe='s3.pipe',
                    mysql_stderr='mysql.mysql_stderr')))

    @property
    def export_snowflake_source_activity(self):
        """Export dimension source from Snowflake.

        Export dimension source input from Snowflake to a gzip file in S3.
        """
        return self.create(
            name='export_snowflake_source_table',
            tasks=runner.Sync(
                db.snowflake_copy_to_s3.fill(
                    s3_path='dim_refresh.staging_path',
                    query='dim_refresh.export_sql',
                    sf_config=StaticParam(SF_CONFIG),
                    aws=StaticParam(self.conf_aws),
                    file_format=StaticParam(
                        self.conf_dim_snowflake_options['file_format']))))

    @property
    def insert_new_dimensions_activity(self):
        """Insert new dimensions.

        Run Snowflake query(s) to insert new dimensions.
        """
        return self.create(
            name='insert_new_dimensions',
            tasks=runner.Sync(
                db.snowflake_execute_query_list.fill(
                    query_list='dim_refresh.insert_sql',
                    sf_config=StaticParam(SF_CONFIG))))

    @property
    def load_snowflake_staging_table_activity(self):
        """Load staging data into Snowflake.

        Truncate the staging table and load new data into it from s3.
        """
        return self.create(
            name='load_snowflake_staging_table',
            retry=5,
            tasks=runner.Sync(
                db.snowflake_execute_query.fill(
                    sf_config=StaticParam(SF_CONFIG),
                    query='dim_refresh.empty_staging_table_sql'),
                db.snowflake_copy_from_s3.fill(
                    sf_config=StaticParam(SF_CONFIG),
                    aws=StaticParam(self.conf_aws),
                    table='dim_refresh.staging_table',
                    s3_path='dim_refresh.staging_path',
                    file_format=StaticParam(
                        self.conf_dim_snowflake_options['file_format']))))

    @property
    def related_dimension_updates_activity(self):
        """Update related dimension tables.

        Run (if any) Snowflake queries to updated related dimension data.
        """
        return self.create(
            name='update_related_dimensions',
            tasks=runner.Sync(
                db.snowflake_execute_query_list.fill(
                    sf_config=StaticParam(SF_CONFIG),
                    query_list='dim_refresh.related_sql')))

    @property
    def update_existing_dimensions_activity(self):
        """Update existing dimensions.

        Run Snowflake query(s) to update existing dimensions.
        """
        return self.create(
            name='update_existing_dimensions',
            tasks=runner.Sync(
                db.snowflake_execute_query_list.fill(
                    sf_config=StaticParam(SF_CONFIG),
                    query_list='dim_refresh.update_sql')))

    @property
    def update_new_dimensions_activity(self):
        """Update new dimensions.

        Run Snowflake query(s) to update existing dimensions.
        """
        return self.create(
            name='update_new_dimensions',
            tasks=runner.Sync(
                db.snowflake_execute_query_list.fill(
                    sf_config=StaticParam(SF_CONFIG),
                    query_list='dim_refresh.update_new_rows_sql')))

    @property
    def delete_old_dimensions_activity(self):
        """Delete old dimensions.

        Run Snowflake query(s) to delete old dimensions.
        """
        return self.create(
            name='delete_old_dimensions',
            tasks=runner.Sync(
                db.snowflake_execute_query_list.fill(
                    sf_config=StaticParam(SF_CONFIG),
                    query_list='dim_refresh.delete_sql')))

    @property
    def notify_dimension_refresh_activity(self):
        """Send sns alert that dimension has been updated."""
        return self.create(
            name='notify_dimension_refresh',
            tasks=runner.Sync(
                notification.record_refresh.fill(
                    namespace='record_refresh',
                    sf_config=StaticParam(SF_CONFIG),
                    dim_type='dim_refresh.dim_type',
                    timestamp='dim_refresh.current_timestamp',
                    insert_count_sql='dim_refresh.insert_count_sql',
                    update_count_sql='dim_refresh.update_count_sql'),
                garcon_sns.sns_publish_message.fill(
                    topic=SNS_CONFIG.get('topic'),
                    message='record_refresh.sns_message',
                    subject='record_refresh.sns_subject')))

    @property
    def downstream_dim_flows_activity(self):
        """Run any dependent dim workflows."""
        return self.create(
            name='downstream_dim_flows',
            tasks=runner.Sync(
                downstream.trigger_workflow.fill(
                    downstream_contexts='downstream_contexts')))