"""Query snowflake for total sales (USD) and total rows; update sales_file."""

from decimal import Decimal
from decimal import ROUND_UP

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_get_eligible.helpers import get_event_from_params
from templates.sales_get_eligible.update_sales_file_metadata \
    import get_distro_sales_file_metadata
from templates.sales_get_eligible.update_sales_file_metadata \
    import get_nr_sales_file_metadata


def update_sales_metadata(dag_run: dict, **kwargs) -> None:
    """Update sales_file record with amount_usd and row_count by querying snowflake.

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

    metadata = _query_metadata_from_snowflake(sales_file_id)
    update_params = _format_update_params(metadata)

    print('UPDATE PARAMS', update_params)

    ows.update_sales_file(sales_file_id, **update_params)


def _format_update_params(metadata: dict) -> dict:
    """Format result of snowflake query.

    Args:
        metadata (dict): result of snowflake query

    Returns:
        new dict of formatted data
    """
    precision = Decimal('0.01')
    raw_amount, row_count = metadata
    rounded_amount = Decimal(raw_amount).quantize(precision, rounding=ROUND_UP)
    return {
        'amount_usd': str(rounded_amount),
        'row_count': row_count
    }


def _query_metadata_from_snowflake(sales_file_id: int) -> dict:
    """Query stmt_db_sales_*_temp for total amount_usd and row_count.

    Args:
        sales_file_id (int): ID of the sales file

    Returns:
        query result as a dict
    """
    sales_file = ows.get_sales_file_details(sales_file_id)
    accounting_period = ows.get_accounting_period_details(
        sales_file['accounting_period_id']
    )
    contract_type = accounting_period['contract_type']

    hook = RoyaltySnowflakeHook(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME)

    sales_file_metadata_query_templates = {
        CONTRACT_TYPES.DISTRIBUTION: get_distro_sales_file_metadata,
        CONTRACT_TYPES.NEIGHBOURING_RIGHTS: get_nr_sales_file_metadata
    }

    snowflake_query = sales_file_metadata_query_templates[contract_type]().render(
        env=config.OWS_ENV,
        sales_file_id=sales_file_id
    )
    result = hook.get_first(snowflake_query)

    print('METADATA', result)

    return result