"""Helper functions for the `adjustment_file_generate` DAG."""

from lib.abacus_event import AbacusEvent
from lib.constants import (
    DAG_ADJUSTMENT_FILE_GENERATE_FLOWTHROUGH_EVENT_NAME,
    DAG_ADJUSTMENT_FILE_GENERATE_TARGET_TYPE,
)
from lib.utils import event, ows


def get_event_from_params(dag_run, **kwargs) -> AbacusEvent:
    """Get the Abacus event from the params passed to task callbacks."""
    abacus_event = event.get_abacus_event(dag_run, **kwargs)

    event.validate_event_for_handler(
        abacus_event,
        target_type=DAG_ADJUSTMENT_FILE_GENERATE_TARGET_TYPE,
        event_names=[DAG_ADJUSTMENT_FILE_GENERATE_FLOWTHROUGH_EVENT_NAME]
    )

    return abacus_event


def get_abacus_state(
    action_name: str,
    statement_period_adjustment_file_id: int
) -> dict:
    """Get the Abacus state for the specified action."""
    abacus_states = ows.get_abacus_states(
        DAG_ADJUSTMENT_FILE_GENERATE_TARGET_TYPE,
        statement_period_adjustment_file_id
    )

    action_state = next(
        (action for action in abacus_states if action['action_name'] == action_name),
        None
    )

    if not action_state:
        raise ValueError(f'{action_name} state not found')

    return action_state