"""Tasks for Amazon DataPulse Ingestion Workflow."""

import datetime

from garcon import task
from garcon_contrib.dynamo_feed_status import garcon_feed_status
from snowflake_connector.etl_connector import SQLLoader

from feed_ingestion.flows.amazon_datapulse import config
from feed_ingestion.tasks import check_status
from feed_ingestion.tasks import s3_tasks
from feed_ingestion.util import sql_templating
from feed_ingestion.util.aws import athena

STOP_RESPONSE = {'stop': True}


@task.decorate(timeout=700)
def bootstrap(activity, date, reload, report_name, partition=None):
    """Bootstrap workflow by injecting initial context from config.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Reporting date (YYYY-MM-DD).
        reload (str or None): If 'True' reload status.
        report_name (str): One of config.reports keys.
        partition (dict or None): Optional partition filter. Keys must exist
            in the report's partitions config; values must match their types.

    Returns:
        dict: Initial context for the workflow.
    """
    assert report_name in config.reports, f'unsupported report "{report_name}"'

    report_config = config.reports[report_name]

    # TODO: validate date and partition report_date
    partition = partition or {}
    if partition:
        # lowercase all partition keys to match report config
        partition = {k.lower(): v for k, v in partition.items()}
        report_partitions = report_config['partitions']
        if not set(partition.keys()) == set(report_partitions):
            raise ValueError(
                f'Invalid partition "{list(partition.keys())}" '
                f'for report "{report_name}". '
                f'Expected: {list(report_partitions.keys())}'
            )
        for key, value in partition.items():
            expected_type = report_partitions[key]
            if not isinstance(value, expected_type):
                raise ValueError(
                    f'Invalid type for partition key "{key}": '
                    f'expected {expected_type.__name__}, '
                    f'got {type(value).__name__}'
                )

        partition_report_date = config.get_partition_date(partition)
        context_date = date
        if (
            partition_report_date is not None
            and context_date != partition_report_date
        ):
            raise ValueError(
                f'Inconsistent dates "{partition_report_date=}" and '
                f'"{context_date=}". '
            )

        # build partition_path ordering by report_partitions keys
        partition_path = '/'.join(
            [f'{key}={partition[key]}'
             for key in report_config['partitions'].keys()]
        )
        archive_path = config.archive_s3_partitioned_path_template.format(
            report_name=report_name,
            partition=partition_path,
        )
    else:
        archive_path = config.archive_s3_full_path_template.format(
            report_name=report_name,
            date=date,
        )

    feed_name = config.get_contextified_feed_name(report_name, partition)
    date_obj = datetime.datetime.strptime(date, config.DATE_FORMAT)
    date = date_obj.strftime(config.DATE_FORMAT)

    if reload == 'True':
        activity.logger.info(
            'Delete status for feed: {} {}'.format(feed_name, date)
        )
        garcon_feed_status.delete_status(feed_name, date)
    elif (
        garcon_feed_status.get_overall_status(feed_name, date)
        == garcon_feed_status.STATUS_INGESTED
    ):
        activity.logger.info('Feed already ingested for {}'.format(date))
        return {
            'stop': True,
            'message': '{} is already ingested for {}'.format(feed_name, date),
        }

    archive_bucket = config.archive_bucket
    s3_dir_path = f's3://{archive_bucket}/{archive_path}'

    staging_raw_database, staging_raw_schema, staging_raw_table = (
        config.get_staging_raw_database_schema(report_name)
    )
    return dict(
        feed_name=feed_name,
        executor_feed_name=config.feed_name,
        report_name=report_name,
        date=date,
        secrets_path=config.secrets_path,
        archive_bucket=archive_bucket,
        archive_path=archive_path,
        s3_dir_path=s3_dir_path,
        athena_source_table=report_config['athena_source_table'],
        staging_raw_table=staging_raw_table,
        sfdb_params={},
        partition=partition,
        kwargs=dict(
            partition=partition,
            staging_raw_database=staging_raw_database,
            staging_raw_schema=staging_raw_schema,
        ),
    )


@task.decorate(timeout=3 * 60 * 60)
@check_status()
def fetch_from_athena(
    activity,
    date,
    feed_name,
    report_name,
    destination_s3_bucket,
    destination_s3_path,
    partition=None,
):
    """Fetch report data from Athena and write PARQUET files to S3.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Reporting date (YYYY-MM-DD).
        feed_name (str): Feed name used for DynamoDB status tracking.
        report_name (str): Report name.
        destination_s3_bucket (str): S3 bucket for PARQUET output.
        destination_s3_path (str): S3 path within the bucket (must end with /).
        partition (dict or None): Optional partition filter applied as
            additional AND conditions in the Athena query.

    Returns:
        dict: source_files_dict with list of PARQUET file metadata.
    """
    if not destination_s3_path.endswith('/'):
        raise ValueError('destination_s3_path should end with /')
    if report_name not in config.reports:
        raise ValueError(f'unsupported report "{report_name}"')
    report_config = config.reports[report_name]

    partition = partition or {}

    sql_loader = SQLLoader(__file__)
    query_name = f'athena_load_{report_name}'
    sql_template = sql_loader.load_query(query_name)
    sql_query, sql_bind_params = sql_templating.render(
        template=sql_template,
        engine='athena',
        params=dict(
            athena_database=config.athena_source_database,
            athena_table=report_config['athena_source_table'],
            date=date,
            partition=partition,
        ),
    )

    athena_output_location = (
        f's3://{destination_s3_bucket}/{destination_s3_path}'
    )

    activity.logger.info(f'Cleaning files at {athena_output_location}')
    s3_tasks.remove_files_from_path(
        activity=activity,
        path=athena_output_location,
        return_deleted_files=False,
    )

    athena.run_query(
        athena_query=sql_query,
        parameters=sql_bind_params,
        athena_temp_database='',
        athena_workgroup=config.athena_workgroup,
        destination_s3_bucket=destination_s3_bucket,
        destination_s3_path=destination_s3_path,
        use_unload_query=True,
    )

    try:
        files = s3_tasks.source_files(
            activity=activity,
            s3_bucket=destination_s3_bucket,
            s3_path=destination_s3_path,
            file_pattern=athena.PARQUET_FILE_PATTERN,
        )
        for file in files['source_files_dict']['files']:
            file['found'] = True
        garcon_feed_status.set_overall_status(
            feed_name, date, garcon_feed_status.STATUS_DOWNLOADED
        )
        activity.logger.info('fetch_from_athena completed')
        return files
    except ValueError:
        error_message = (
            f'Empty dataset. No parquet '
            f'data files were found in {destination_s3_path}'
        )
        activity.logger.error(error_message)
        garcon_feed_status.set_overall_status(
            feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE
        )
        raise ValueError(error_message)