"""Component for loading source files directly to staging_raw via stage."""

from datetime import datetime

from garcon import task
from garcon.param import StaticParam
from snowflake_connector.etl_connector import SnowflakeSQLExecutor, SQLLoader

from feed_ingestion.common.staging_raw_sf.snowflake_stage_loader import \
    StageLoader
from feed_ingestion.conf.config import merge_configs
from feed_ingestion.flows import base
from feed_ingestion.flows.helpers import get_sf_config
from feed_ingestion.util import task_status


class SnapchatSL(StageLoader):
    """StageLoader class."""

    sql_loader = SQLLoader(__file__)

    def load_staging_raw_table(self, staging_raw_table, source_files_dict,
                               date, stage_name, **kwargs):
        """Load the temp_staging_raw data to the feed's staging_raw table.

        Args:
            staging_raw_table (str): A table name in Snowflake.
            source_files_dict (dict): A dict with source files metadata.
            date (str): Date of the data being process (YYYY-MM-DD).
            stage_name (str): Name of Snowflake stage containing source files.
        """
        ingestion_time = datetime.now()

        query_name = kwargs.get('query_name')

        for file_dict in source_files_dict['files']:
            self.resolve_sql_loader_and_execute(
                query_name, params=dict(
                    db=self.executor.sf_config['db'],
                    schema=self.executor.sf_config['schema'],
                    stage=stage_name,
                    staging_raw_table=staging_raw_table,
                    file_name=file_dict['file_name'].lstrip('/'),
                    file_size=file_dict['file_size'],
                    download_date=date,
                    ingestion_time=ingestion_time,
                    **kwargs))

    @classmethod
    def load_activity(
            cls, feed_name, requirements, sql_loader=None,
            executor_class=SnowflakeSQLExecutor, secrets_path=None,
            task_timeout=7200):
        """Load from stage activity.

        Args:
            feed_name (str): name of the feed.
            requirements (dict): activity requirements dict.
            sql_loader (SQLLoader): sql loader instance.
            executor_class: Snowflake executor class.
            secrets_path (str): Secrets manager path of the flow.
        """
        @task.decorate(timeout=task_timeout)
        def load_task(
                activity, date, s3_dir_path, staging_raw_table_name,
                source_files_dict, sfdb_params, aws,
                report=None,
                skip_corrupted_rows=False):
            """Copy data from Snowflake stage to staging_raw table.

            Args:
                activity (ActivityWorker): The activity worker.
                date (str): Reporting date (YYYY-MM-DD).
                s3_dir_path (str): Name of the feed to get executor class.
                staging_raw_table_name (str): name of staging raw table.
                source_files_dict (dict): Name of the temp staging_raw table.
                report (str): report name set query_name
                    for load_staging_raw task. When defined then
                    query_name == 'load_staging_raw_{report}.sql'.
                    if not specified then query_name == 'load_staging_raw.sql'
                sfdb_params (dict): Dict with params to optionally override
                    default ones (Snowflake db and schema name).
                aws (dict): aws credentials.
                skip_corrupted_rows (bool): If True, add ON_ERROR=CONTINUE.
            """
            if task_status.is_completed_task(
                    feed_name + '_' + requirements.get('report'), date,
                    'load_staging_raw_table'):
                activity.logger.info(
                    'task: {task_id} date: {date} status: COMPLETE'.format(
                        task_id='load_staging_raw_table', date=date))
                return

            activity.logger.info('Loading staging raw table: %s', date)

            sf_config = get_sf_config(secrets_path)
            sf_config_custom = merge_configs(sf_config, sfdb_params)
            with executor_class(sf_config_custom) as executor:
                stage_loader = cls(executor, sql_loader)
                stage_loader.clean_staging_raw_table(
                    staging_raw_table_name, date)

                stage_name = '{feed_name}_stage_{date:%Y%m%d}'.format(
                    feed_name=feed_name,
                    date=datetime.strptime(date, '%Y-%m-%d'))

                stage_loader.create_stage(stage_name, s3_dir_path, aws)

                args = [
                    staging_raw_table_name, source_files_dict, date,
                    stage_name]
                kwargs = {}
                if skip_corrupted_rows:
                    kwargs['skip_corrupted_rows'] = True
                if report:
                    kwargs['query_name'] = f'load_staging_raw_{report}'
                stage_loader.load_staging_raw_table(*args, **kwargs)

                stage_loader.drop_stage(stage_name)

        return property(
            lambda flow_self: flow_self.create(
                name='load_staging_raw_table_from_stage',
                tasks=base.SyncRunner(
                    load_task.fill(
                        aws=StaticParam(flow_self.conf_aws),
                        **requirements)
                )
            )
        )