"""Task to export sales from snowflake to s3 as a parquet file."""

from hooks.royalty_snowflake_hook import RoyaltySnowflakeHook
from lib import config
from lib.constants import CONTRACT_TYPES
from lib.utils import ows
from tasks.sales_approve import helpers
from templates.sales_approve.snowflake_export_sales\
    import export_approved_distro_sales_to_s3
from templates.sales_approve.snowflake_export_sales\
    import export_approved_nr_sales_to_s3


def export_sales_from_snowflake(dag_run: dict, **kwargs) -> None:
    """Copy data from stmt_db_sales_* to s3 and update sales_file with s3 path.

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

    sales_file, accounting_period = helpers.get_event_records(event)
    s3_location = helpers.build_sales_file_main_url(accounting_period, sales_file)

    contract_type = accounting_period['contract_type']
    approved_sales_templates = {
        CONTRACT_TYPES.DISTRIBUTION: export_approved_distro_sales_to_s3,
        CONTRACT_TYPES.NEIGHBOURING_RIGHTS: export_approved_nr_sales_to_s3
    }

    hook = RoyaltySnowflakeHook(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME)
    export_statement = approved_sales_templates[contract_type]().render(
        env=config.OWS_ENV,
        s3_path=s3_location.key,
        sales_file_id=sales_file_id,
        stage=config.ABACUS_PARQUET_STAGE
    )
    hook.run(export_statement, autocommit=True)

    ows.update_sales_file(sales_file_id, **{'main_url': s3_location.url})