"""Tasks for s32sf workflow."""

import datetime
import difflib
import json

from garcon import task

from garcon_contrib.dynamo_feed_status import garcon_feed_status
from garcon_contrib.snowflake import garcon_snowflake

from snowflake_etl.conf import config as cfg
from snowflake_etl.conf.config import merge_configs
from snowflake_etl.conf.config import SF_CREDENTIALS
from snowflake_etl.flows.s32sf import helpers
from snowflake_etl.util import schema_map

SOURCES_CONF = helpers.SourcesConf()


@task.decorate(timeout=1000)
def bootstrap(
        activity, table, s3_path_data, s3_path_schema, load_strategy,
        sfdb_params=None, file_format=None, date=None, feed_name=None):
    """Bootstrap S3 to Snowflake flow.

    Checking required params are provided

    Args:
        activity (ActivityWorker): The activity worker.
        sfdb_params (dict): a dict sf override params (no creds).
        table (str): Name of the table on Snowflake.
        s3_path_data (str): Path to data on S3.
        s3_path_schema (str): Path to JSON file for schema.
        load_strategy (str): Strategy of loading Snowflake
            (snapshot or incremental).
        file_format (list): Optional FILE_FORMAT parameter for Snowflake
            ingestion.
        date (str): Optional a date of ingestion into Snowflake.
        feed_name (str): Optional feed name, provided by a calling flow.

    Returns:
        dict: Context dictionary.
    Raises:
        Exception: If a value is missing from the context, an error is thrown.
    """
    activity.logger.info('bootstrapping S3-to-Snowflake workflow')

    sfdb_params = sfdb_params or {}
    sfdb_params = merge_configs(cfg.SF_PARAMS, sfdb_params)
    db, schema = helpers.extract_db_and_schema(sfdb_params)

    assert db, 'db is required for the flow to work.'
    assert schema, 'schema is required for the flow to work.'
    assert table, 'table is required for the flow to work.'
    assert s3_path_data, 's3_path_data is required for the flow to work.'
    assert s3_path_schema, 's3_path_schema is required for the flow to work.'
    assert load_strategy in cfg.SF_LOAD_STRATEGIES, 'invalid load_strategy'

    return dict(
        sfdb_params=sfdb_params,
        table=table,
        s3_path_data=s3_path_data,
        s3_path_schema=s3_path_schema,
        load_strategy=load_strategy,
        file_format=file_format,
        date=date,  # Currently is set in s32sf etl flow only
        feed_name=feed_name
    )


# FIXME this should be a lot more generic
@task.decorate(timeout=2000)
def source_schema(activity, s3_path_schema):
    """Extract table schema from a json file on S3.

    Args:
        activity (ActivityWorker): The activity worker.
        s3_path_schema (str): S3 path to JSON schema file.

    Returns:
        dict: Context with source table schema translated to Snowflake SQL.
    """
    activity.logger.info(
        'fetching external schema from {schema}'.format(schema=s3_path_schema))

    schema_json = json.loads(
        helpers.read_json_schema_on_s3(s3_path_schema, activity.logger))

    return dict(source_schema=schema_json)


@task.decorate(timeout=2000)
def create_destination_table(
        activity, sfdb_params, table, source_schema, source_schema_name=None,
        db_type='redshift'):
    """Create Snowflake table if it doesn't already exist from source schema.

    Args:
        activity (ActivityWorker): The activity worker.
        sfdb_params (dict): a dict sf connection parameters (no creds).
        table (str): Table name.
        source_schema (dict): Source schema.
        db_type (str): Type of a source database ('mysql' or 'redshift').
    """
    db, dest_schema = helpers.extract_db_and_schema(sfdb_params)
    sf_config = merge_configs(sfdb_params, SF_CREDENTIALS)

    activity.logger.info(
        'checking if {db}.{schema}.{table} '
        'already exists in snowflake'.format(
            db=db, schema=dest_schema, table=table))

    if not (garcon_snowflake.table_exists(table, sf_config=sf_config)):

        activity.logger.info('creating {db}.{schema}.{table}'.format(
            db=db, schema=dest_schema, table=table))

        # convert source schema
        if db_type == 'mysql':
            table_schema = schema_map.SQLTable.from_json(**source_schema)
            sf_sql = table_schema.sql_snowflake(db_type)

        # if no source_schema default to not transient
        if source_schema_name is None:
            transient = ''
        else:
            transient = SOURCES_CONF.is_table_transient(
                table, source_schema_name)

        # prepare create table sql
        sql = cfg.SF_DEFAULT_QUERIES['create_dst_table'].format(
            db=db,
            schema=dest_schema,
            transient=transient,
            table=table,
            columns=', '.join(sf_sql))

        activity.logger.info('>> Executing query {sql}'.format(sql=sql))

        # create table
        garcon_snowflake.execute_with_py_conn(sql, sf_config=sf_config)
    else:
        activity.logger.info(
            'table already exists {db}.{schema}.{table}'.format(
                db=db, schema=dest_schema, table=table))


@task.decorate(timeout=2000)
def target_schema(activity, sfdb_params, table):
    """Extract schema from Snowflake.

    Args:
        activity (ActivityWorker): The activity worker.
        sfdb_params (dict): a dict sf connection parameters (no creds).
        table (str): Table name in Snowflake.

    Returns:
        dict: Context with target Snowflake schema.
    """
    db, schema = helpers.extract_db_and_schema(sfdb_params)
    sf_config = merge_configs(sfdb_params, SF_CREDENTIALS)

    sql = cfg.SF_DEFAULT_QUERIES['describe_table'].format(
        db=db, schema=schema, table=table)

    activity.logger.info(
        'fetching schema for {db}.{schema}.{table}'
        '\n>> Executing query {sql}'.format(
            db=db, schema=schema, table=table, sql=sql))

    cols = garcon_snowflake.execute_with_py_conn(
        sql, garcon_snowflake.FetchEnum.ALL, sf_config=sf_config)['results']

    # FIXME a little hacky
    target_schema = tuple('{name} {type} {nullable}'.format(
        name=row[0].replace('`', ''),
        type=row[1],
        nullable='' if row[3] == 'Y' else 'NOT NULL').strip() for row in cols)

    return dict(target_schema=target_schema)


@task.decorate(timeout=2000)
def validate_compatibility(
        activity, source_schema, target_schema, strict='true',
        db_type='redshift'):
    """Throw Assertion Error if type or order mismatches found.

    Args:
        activity (ActivityWorker): The activity worker.
        source_schema (dict): Source schema in raw dict format.
        target_schema (tuple): Target schema in a tuple of SQL.
        strict (str): Validate and assert or just warn (default True).
        db_type (str): Type of a source database ('mysql' or 'redshift').

    Exception:
        AssertionError: Indicates a schema mismatch.
    """
    activity.logger.info(
        'validating compatibility of source and target schemas')

    # FIXME convert source schema to target format
    if db_type == 'mysql':
        table_schema = schema_map.SQLTable.from_json(**source_schema)
        sf_sql = table_schema.sql_snowflake(db_type)

    activity.logger.info('Source SQL: {}'.format(sf_sql))
    activity.logger.info('Target Schema: {}'.format(target_schema))

    diff = tuple(difflib.ndiff(sf_sql, target_schema))

    activity.logger.info('schema diff:\n %s', '\n'.join(diff))

    # the value comes in the string form (str(True))
    if str(strict).lower() == 'false':
        activity.logger.warn('lax validation - continuing')
        return

    activity.logger.info('Diff {} vs {}'.format(len(diff), len(sf_sql)))
    activity.logger.info('Diff: {}'.format(diff))
    assert len(diff) == len(sf_sql), 'Schema mismatches: {}'.format(diff)


@task.decorate(timeout=2000)
def create_staging_table(
        activity, sfdb_params, table, source_schema_name=None,
        db_type='redshift'):
    """Build create replace temp table SQL.

    Args:
        activity (ActivityWorker): The activity worker.
        sfdb_params (dict): a dict sf connection parameters (no creds).
        table (str): Name of the Snowflake table.
        db_type (str): Type of a source database ('mysql' or 'redshift').

    Return:
        str: SQL for create replace Snowflake table.
    """
    db, dest_schema = helpers.extract_db_and_schema(sfdb_params)
    sf_config = merge_configs(sfdb_params, SF_CREDENTIALS)

    stage_table = 'stg_{}'.format(table)
    if garcon_snowflake.table_exists(stage_table, sf_config=sf_config):
        activity.logger.info('{} table already exists in {}'.format(
            stage_table, dest_schema))
        return

    if db_type == 'mysql':
        comment = 'Automated SQL sync | {date}'.format(
            date=datetime.datetime.now().strftime('%Y/%m/%d %H:%M:%S'))

    # if no source_schema default to not transient
    if source_schema_name is None:
        transient = ''
    else:
        transient = SOURCES_CONF.is_table_transient(
            table, source_schema_name)

    sql = cfg.SF_DEFAULT_QUERIES['create_stg_table'].format(
        db=db,
        schema=dest_schema,
        transient=transient,
        table=table,
        comment=comment)

    activity.logger.info(
        'creating staging table: {db}.{schema}.{table}'
        '\n>> Executing query {sql}'.format(
            db=db, schema=dest_schema, table=table, sql=sql))

    garcon_snowflake.execute_with_py_conn(sql, sf_config=sf_config)


# FIXME a lot of these can be genericized


@task.decorate(timeout=50400)
def load_staging_table(
        activity, sfdb_params, table, s3_path_data, file_format=None,
        db_type='redshift'):
    """Load external S3 data.

    Args:
        activity (ActivityWorker): The activity worker.
        sfdb_params (dict): a dict sf connection parameters (no creds).
        table (str): Name of the Snowflake table.
        s3_path_data (str): Full S3 path to data source.
        file_format (list): Optional FILE_FORMAT parameter for Snowflake
            ingestion.
        db_type (str): Type of a source database ('mysql' or 'redshift').

    Return:
        str: SQL for create replace Snowflake table.
    """
    db, schema = helpers.extract_db_and_schema(sfdb_params)
    sf_config = merge_configs(sfdb_params, SF_CREDENTIALS)

    aws_conf = cfg.getconf('aws')['aws']

    if file_format:
        helpers.validate_file_format(file_format)
        file_format = ' '.join(file_format)
    else:
        if db_type == 'redshift' or db_type is None:
            file_format = cfg.SF_DEFAULT_FILE_FORMAT
        elif db_type == 'mysql':
            file_format = cfg.SF_MYSQL_STDOUT_FILE_FORMAT

    sql = cfg.SF_DEFAULT_QUERIES['truncate_stg_table'].format(
        db=db,
        schema=schema,
        table=table)
    garcon_snowflake.execute_with_py_conn(sql, sf_config=sf_config)

    activity.logger.info(
        'truncating staging table: {db}.{schema}.stg_{table}'.format(
            db=db, schema=schema, table=table))

    sql = cfg.SF_DEFAULT_QUERIES['copy_into_stg_table'].format(
        db=db,
        schema=schema,
        table=table,
        s3_path_data=s3_path_data,
        aws_key_id=aws_conf['access_key'],
        aws_secret_key=aws_conf['access_secret'],
        file_format=file_format)

    activity.logger.info(
        'loading staging table: {db}.{schema}.{table}'
        '\n>> Executing query {sql}'.format(
            db=db, schema=schema, table=table, sql=sql))

    garcon_snowflake.execute_with_py_conn(sql, sf_config=sf_config)


@task.decorate(timeout=600)
def swap_snowflake_tables(activity, sfdb_params, table, load_strategy):
    """Swap staging and destination Snowflake tables.

    Args:
        activity (ActivityWorker): The activity worker.
        sfdb_params (dict): a dict sf connection parameters (no creds).
        table (str): Name of the Snowflake table.
        load_strategy (str): Strategy of loading the Snowflake table.
    """
    # skip if not snapshot (full) load
    if not load_strategy == 'snapshot':
        activity.logger.info(
            'Swap staging and permanent SF tables: skipping for strategy: '
            '{}'.format(load_strategy))
        return

    db, schema = helpers.extract_db_and_schema(sfdb_params)
    sf_config = merge_configs(sfdb_params, SF_CREDENTIALS)

    sql = cfg.SF_DEFAULT_QUERIES['swap_stg_and_dst_table'].format(
        db=db, schema=schema, table=table)

    activity.logger.info(
        'Swapping destination table {db}.{schema}.{table} with staging table '
        '{db}.{schema}.stg_{table}\n>> Executing query {sql}'.format(
            db=db, schema=schema, table=table, sql=sql))

    garcon_snowflake.execute_with_py_conn(sql, sf_config=sf_config)


@task.decorate(timeout=2000)
def drop_staging_table(activity, sfdb_params, table, load_strategy):
    """Drop old table after swap.

    It's only required with a snapshot strategy, because there is no reason
    to store transient tables, which are just the snapshots of the source
    tables, even for debugging purposes.

    Args:
        activity (ActivityWorker): The activity worker.
        sfdb_params (dict): a dict sf connection parameters (no creds).
        table (str): Name of the Snowflake table.
        load_strategy (str): Strategy of loading the Snowflake table.
    """
    # skip if not snapshot (full) load
    if not load_strategy == 'snapshot':
        activity.logger.info(
            'drop_staging_table: skipping for strategy: {}'.format(
                load_strategy))
        return

    db, schema = helpers.extract_db_and_schema(sfdb_params)
    sf_config = merge_configs(sfdb_params, SF_CREDENTIALS)

    sql = cfg.SF_DEFAULT_QUERIES['drop_stg_table'].format(
        db=db, schema=schema, table=table)

    activity.logger.info(
        'dropping old staging table: {db}.{schema}.{table}'
        '\n>> Executing query {sql}'.format(
            db=db, schema=schema, table=table, sql=sql))

    garcon_snowflake.execute_with_py_conn(sql, sf_config=sf_config)


@task.decorate(timeout=2000)
def insert_into_destination_table(activity, sfdb_params, table, load_strategy):
    """Insert the data from staging table into destination table.

    Args:
        activity (ActivityWorker): The activity worker.
        sfdb_params (dict): a dict sf connection parameters (no creds).
        table (str): Name of the Snowflake table.
        load_strategy (str): Strategy of loading to the Snowflake table.
    """
    # skip if not incremental (partial) load
    if not load_strategy == 'incremental':
        activity.logger.info(
            'insert_into_destination_table: skipping for strategy: {}'.format(
                load_strategy))
        return

    db, schema = helpers.extract_db_and_schema(sfdb_params)
    sf_config = merge_configs(sfdb_params, SF_CREDENTIALS)

    sql = cfg.SF_DEFAULT_QUERIES['insert_into_dst_table'].format(
        db=db, schema=schema, table=table)

    activity.logger.info(
        'inserting into destination table: {db}.{schema}.{table}'
        '\n>> Executing query {sql}'.format(
            db=db, schema=schema, table=table, sql=sql))

    garcon_snowflake.execute_with_py_conn(sql, sf_config=sf_config)


@task.decorate(timeout=1000)
def set_ingestion_status(activity, date, feed_name, status):
    """Explicitly set the overall status of a feed.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date (YYYY-MM-DD).
        feed_name (str): Name of the feed.
        status (str): Status constant in util.feed_status.
    """
    if date is None or feed_name is None:
        activity.logger.info(
            'Reporting date is not provided. Skipping status setting '
            'in DynamoDB')
        return

    garcon_feed_status.set_overall_status(feed_name, date, status)
    activity.logger.info(
        'Setting status for feed: {feed_name} '
        'with date: {date} '
        'to: {status}, via a task'.format(
            feed_name=feed_name, date=date, status=status))