"""Tasks for geocoding workflow."""

from garcon import task
import geocoder
import smart_open
try:
    from snowflake_connector.etl_connector import SnowflakeSQLExecutor
except ImportError:
    # in order to decouple geocoding which has different dependencies from
    # other workflows
    SnowflakeSQLExecutor = None

from dim_refresh_etl.conf.config import merge_configs
from dim_refresh_etl.conf.config import SF_CONFIG
from dim_refresh_etl.flows.geocoding import settings


def _write_lat_long_to_s3_file(file_object, country_code, zip_code):
    """Get country and zip codes for requesting the google Geocoding API.

    Args:
        file_object (smart_open file object): smart_open file object.
        country_code (str): 2 letter country code from DB.
        zip_code (str): postal code.
    """
    lat_long = geocoder.google(
        '', components='country:{}|postal_code:{}'.format(
            country_code, zip_code))
    row_placeholder = '{}\t{}\t{}\t{}\n'
    if lat_long.latlng:
        file_object.write(
            row_placeholder.format(
                zip_code, country_code, lat_long.latlng[0],
                lat_long.latlng[1]))
    else:
        file_object.write(
            row_placeholder.format(
                zip_code, country_code, '', ''))


def _get_zip_country_code_sql(country_code):
    """Get and modify SQL that gets zip and country code.

    Args:
        sql (str): SQL statement.

    Returns:
        str: SQL statement.
    """
    sql = settings.SELECT_SQL
    if country_code:
        sql += " AND country_code = '{country_code}'"
    sql += (
        " AND requested = 'N' "
        'ORDER BY date_created DESC LIMIT {}').format(settings.REQUEST_LIMIT)
    return sql.format(country_code=country_code)


@task.decorate(timeout=18000)
def write_country_zip_code_lat_long_map(
        activity, sfdb_params, country_code=None):
    """Generate country-zipcode to latitude/longitude mapping file.

    Produce a mapping data file on s3 to map latitude and longitude to zip
    and country code pairs.

    Args:
        activity (ActivityWorker): the activity worker.
        country_code (str): optional country code. If specified, that
            country will be prioritized.
        sfdb_params (dict): Snowflake connection params.
    """
    sf_config = merge_configs(SF_CONFIG, sfdb_params)

    with SnowflakeSQLExecutor(sf_config) as executor:
        with executor.get_cursor() as cursor:
            sql, _ = executor.validator.format_identifiers(
                _get_zip_country_code_sql(country_code), {
                    'db': sf_config['db'],
                    'schema': sf_config['schema']})
            activity.logger.info('Running sql: {}'.format(sql))
            cursor.execute(sql)

            with smart_open.smart_open(settings.TEMP_S3_PATH, 'w') as output:
                for country_code, zip_code in cursor:
                    activity.logger.info(
                        'Writing latlong for: {} [{}]'.format(
                            country_code, zip_code))
                    _write_lat_long_to_s3_file(
                        output, country_code, zip_code)


@task.decorate(timeout=7200)
def update_lat_long_on_dim_zip(activity, sfdb_params):
    """Update latitude/longitude on dim_zip table with the mapping file.

    Args:
        activity (ActivityWorker): the activity worker.
        sfdb_params (dict): Snowflake connection params.
    """
    sf_config = merge_configs(SF_CONFIG, sfdb_params)
    with SnowflakeSQLExecutor(sf_config) as executor:
        sql, _ = executor.validator.format_identifiers(
            settings.CREATE_TEMP_TABLE_SQL, {
                    'db': sf_config['db'],
                    'schema': sf_config['schema'],
                    'table_name': settings.TEMP_TABLE_NAME
                })
        activity.logger.info('Creating temp table for dim_zip_mapping.')
        executor.execute(sql)

        sql, _ = executor.validator.format_identifiers(
            settings.LOAD_TEMP_TABLE_SQL, {
                'db': sf_config['db'],
                'schema': sf_config['schema'],
                'table_name': settings.TEMP_TABLE_NAME,
            })
        activity.logger.info('Loading temp table for dim_zip_mapping.')
        executor.execute(sql, {
            's3_path': settings.TEMP_S3_PATH,
            'aws_key_id': settings.AWS_ACCESS_KEY_ID,
            'aws_secret_key': settings.AWS_SECRET_ACCESS_KEY
        })

        sql, _ = executor.validator.format_identifiers(
            settings.UPDATE_SQL, {
                'db': sf_config['db'],
                'schema': sf_config['schema'],
                'table_name': settings.TEMP_TABLE_NAME
            })
        activity.logger.info('Updating dim_zip table.')
        executor.execute(sql)