"""Task to insert flattened NR contract data into denormalized table.

Inserts into one of two contract denormalized tables:
    * CONTRACT_DENORMALIZED_NR
    * CONTRACT_DENORMALIZED_NR_CONTRIB_ONLY
"""

from hooks.royalty_snowflake_hook import RoyaltySnowflakeHook
from lib import config
from tasks.accounting_run_calculate_nr import helpers
from templates.accounting_run_calculate.snowflake_flatten_contract_json_nr import \
    insert_flat_nr_contract_contributor_only_data_template
from templates.accounting_run_calculate.snowflake_flatten_contract_json_nr import \
    insert_flat_nr_contract_data_template


def insert_flat_contract_task(
    dag_run: dict,
    is_contributor_only: bool = False,
    **kwargs
) -> None:
    """Use the snowflake template to insert flat NR contract data into denorm table.

    Args:
        dag_run (dict): config of the DAG this task belongs to
        is_contributor_only (bool): whether to query contributor-only data
        kwargs (dict): any other optional arguments
    """
    event = helpers.get_event_from_params(dag_run, **kwargs)
    accounting_run_id = event.target_id

    hook = RoyaltySnowflakeHook(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME)

    if is_contributor_only:
        insert_statement = insert_flat_nr_contract_contributor_only_data_template() \
            .render(
                accounting_run_id=accounting_run_id,
                schema=config.OWS_ENV
        )
    else:
        insert_statement = insert_flat_nr_contract_data_template().render(
            accounting_run_id=accounting_run_id,
            schema=config.OWS_ENV
        )

    hook.run(insert_statement, autocommit=True)