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

from datetime import datetime

from garcon import task
from snowflake_connector.etl_connector import SnowflakeSQLExecutor

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.tasks import check_status


class KugouSL(StageLoader):
    """StageLoader class."""

    def create_stage(self, stage_name, s3_dir_path, aws, **kwargs):
        """Create Snowflake stage.

        Args:
            stage_name (str): Name of Snowflake stage containing source files.
            s3_dir_path (str): s3 path to directory containing source files.
            aws (dict): aws credentials.
        """
        aws_params = self.get_aws_params()
        query = 'create_snowflake_stage'
        self.resolve_sql_loader_and_execute(
            query, params=dict(
                db=self.executor.sf_config['db'],
                schema=self.executor.sf_config['schema'],
                stage=stage_name,
                s3_dir_path=s3_dir_path,
                **aws_params
            )
        )

    def load_staging_raw_table(
            self, staging_raw_table, source_files_dict, date, stage_name,
            skip_corrupted_rows, **kwargs):
        """Load the 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.
            skip_corrupted_rows (str): If 'True', skip the corrupted rows.
        """
        ingestion_time = datetime.now()
        query_name = 'load_staging_raw'
        for file_dict in source_files_dict['files']:
            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'],
                file_size=file_dict['file_size'],
                download_date=date,
                ingestion_time=ingestion_time,
                **kwargs)
            if skip_corrupted_rows:
                params.update({'on_error': 'CONTINUE'})
            else:
                params.update({'on_error': 'ABORT_STATEMENT'})
            self.resolve_sql_loader_and_execute(
                query_name, params=params)

    @classmethod
    def load_activity(
            cls, feed_name, requirements, sql_loader=None,
            executor_class=SnowflakeSQLExecutor, secrets_path=None):
        """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=7200)
        @check_status(task_id='load_staging_raw_table')
        def load_task(
                activity, feed_name, date, s3_dir_path,
                staging_raw_table_name,
                source_files_dict, sfdb_params,
                skip_corrupted_rows=False, licensor=None):
            """Copy data from Snowflake stage to staging_raw table.

            Args:
                activity (ActivityWorker): The activity worker.
                feed_name (str): name of the feed.
                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): A dict with source files metadata.
                sfdb_params (dict): Dict with params to optionally override
                    default ones (Snowflake db and schema name).
                skip_corrupted_rows (bool): If True, add ON_ERROR=CONTINUE.
                licensor (str): Optional licensor name.
            """
            activity.logger.info('Loading staging raw table: %s', date)

            # Resolve the date-versioned queries sub folder (e.g.
            # queries/2026-07-01/) for the reporting date. Falls back to the
            # root queries/ folder for dates before any versioned folder.
            if sql_loader:
                sql_loader.date = date
                sql_loader.folder_version = sql_loader._get_sub_folder(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)
                kwargs = {
                    'skip_corrupted_rows': skip_corrupted_rows,
                    'licensor': licensor
                }

                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, None, **kwargs)

                stage_loader.clean_staging_raw_table(
                    staging_raw_table_name, date, **kwargs)

                args = [
                    staging_raw_table_name, source_files_dict, date,
                    stage_name]
                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(**requirements)
                )
            )
        )