"""Task to load sales from s3 to snowflake staging tables."""

from hooks.royalty_snowflake_hook import RoyaltySnowflakeHook
from lib import config
from lib.utils.event import get_abacus_event
from templates.sales_ingest.snowflake_load_sales\
    import insert_sales_to_distro_staging_table
from templates.sales_ingest.snowflake_load_sales\
    import insert_sales_to_nr_staging_table


def load_sales_to_staging_table_task(
    dag_run: dict,
    **kwargs
) -> None:
    """Load sales from S3 files to snowflake staging tables.

    Args:
        dag_run (dict): config of the DAG this task belongs to
        kwargs (dict): any other optional arguments
    """
    event = get_abacus_event(dag_run, **kwargs)
    batch_id = event.target_id

    task_instance = kwargs.get('task_instance')

    sales_ingest_type = task_instance.xcom_pull(
        task_ids='determine_sales_ingest_type',
        key='sales_ingest_type'
    )

    if sales_ingest_type is None:
        raise Exception('Unable to get the type for the sales insertion.')

    print(f'Inserting sales data from S3 files to {sales_ingest_type} staging table for batch {batch_id}')  # noqa: E501

    insertion_template = {
        'distro': insert_sales_to_distro_staging_table,
        'nr': insert_sales_to_nr_staging_table
    }

    hook = RoyaltySnowflakeHook(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME)
    insert_statement = insertion_template[sales_ingest_type]().render(
        batch_id=batch_id,
        schema=config.OWS_ENV.upper(),
        stage=config.ABACUS_TSV_STAGE,
        s3_path=f'extract-sales/{sales_ingest_type}/{batch_id}/parts/'
    )

    try:
        print(f'Executing insert statement: {insert_statement}')
        hook.run(insert_statement, autocommit=True)
    except Exception as e:
        print(f'Error executing insert statement: {e}')
        raise