"""Check Processor."""

from src.connectors.ows_account import (
    account_payee_dataloader,
    account_payment_term_dataloader,
    bulk_get_payment_holds,
    get_accounts,
)
from src.connectors.ows_state import (
    bulk_query_abacus_states,
    StateStatus,
)
from src.constants import (
    ACCOUNT_PAYEE_TABLE,
    COUNTRY_USA,
    TAX_ELIGIBILITY_STATE,
)
from src.models import (
    AbacusState,
    Account,
    AccountPayee,
    AccountPaymentHold,
    AccountPaymentTerm,
    Event,
    PaymentGroup,
)
from src.processors.base.payment_batch_processor import PaymentBatchProcessor
from src.processors.base.processor import Processor
from src.processors.models import EligibleAccountLevelData
from src.utils import fetch_all_account_tax_info_entries


class CheckProcessor(Processor):
    """Processor to determine eligiblity and calculate worksheets for all accounts paid by check."""

    def __init__(self, payment_group: PaymentGroup, abacus_event: Event) -> None:
        """Init."""
        super().__init__(payment_group, abacus_event)
        self._payment_schedules = payment_group.group_criteria.get(
            'payment_schedules', []
        )
        self._reference_payment_type_id = payment_group.group_criteria.get(
            'reference_payment_type_id'
        )
        self._agreement_type_ids = payment_group.group_criteria.get(
            'reference_agreement_types'
        )

    def _set_total_count(self) -> None:
        """Set total count."""
        self._total_count = get_accounts(
            limit=0,
            offset=0,
            reference_payment_type_id=self._reference_payment_type_id,
            agreement_type_ids=self._agreement_type_ids,
        ).total_count

    def _is_account_eligible(
        self,
        account: Account,
        account_payment_term: AccountPaymentTerm,
        account_payee: AccountPayee,
        payment_holds_by_account_id: dict[int, AccountPaymentHold],
        tax_eligibility_states_by_payee_id: dict[int, AbacusState],
    ) -> bool:
        """Determine if given check account is payment eligible."""
        payment_schedule = account_payment_term.payment_schedule
        if (
            len(self._payment_schedules)
            and payment_schedule not in self._payment_schedules
        ):
            # account does not match payment schedule criteria
            return False

        payment_hold = payment_holds_by_account_id.get(account.account_id)
        if payment_hold and payment_hold.is_on_hold:
            # account is on hold and therefore cannot be paid
            return False

        tax_eligibility_state = tax_eligibility_states_by_payee_id.get(
            account_payee.account_payee_id
        )
        if (
            not tax_eligibility_state
            or tax_eligibility_state.action_status != StateStatus.COMPLETE
        ):
            # account is not tax eligible and cannot be paid
            return False

        return True

    def _get_eligible_account_batch_data(
        self, offset: int = 0
    ) -> list[EligibleAccountLevelData]:
        """Get the batch of accounts (alongside necessary account level info) to process for check payments."""
        accounts_resp = get_accounts(
            limit=self._limit,
            offset=offset,
            reference_payment_type_id=self._reference_payment_type_id,
            agreement_type_ids=self._agreement_type_ids,
        )
        accounts_items = accounts_resp.items

        # fetch account payment term data to determine eligibility and augment currency code
        account_ids = [account.account_id for account in accounts_items]
        payment_terms_by_account_id = account_payment_term_dataloader(account_ids)
        account_tax_info = fetch_all_account_tax_info_entries(account_ids)
        account_tax_info_by_account_id = {
            item.account_id: item for item in account_tax_info
        }
        account_payees_by_account_id = account_payee_dataloader(account_ids=account_ids)

        payment_holds_resp = bulk_get_payment_holds(account_ids)
        payment_holds_by_account_id = {
            hold.account_id: hold for hold in payment_holds_resp.items
        }

        account_payee_ids = [
            payee.account_payee_id for payee in account_payees_by_account_id.values()
        ]
        tax_eligibility_resp = bulk_query_abacus_states(
            action_name=TAX_ELIGIBILITY_STATE,
            parent_table_name=ACCOUNT_PAYEE_TABLE,
            parent_table_ids=account_payee_ids,
        )
        tax_eligibility_states_by_payee_id = {
            state.parent_table_id: state for state in tax_eligibility_resp.items
        }

        eligible_account_batch = []
        # filter out ineligible accounts
        for account in accounts_items:
            account_id = account.account_id
            payment_term = payment_terms_by_account_id[account_id]
            tax_info = account_tax_info_by_account_id[account_id]
            account_payee = account_payees_by_account_id[account_id]

            if self._is_account_eligible(
                account,
                payment_term,
                account_payee,
                payment_holds_by_account_id,
                tax_eligibility_states_by_payee_id,
            ):
                formatted_account = EligibleAccountLevelData(
                    account_id=account_id,
                    country_of_tax_residence=tax_info.country_of_tax_residence,
                    country_of_tax_policy=COUNTRY_USA,  # all check payments are USA
                    currency_code=payment_term.currency_code,
                )
                eligible_account_batch.append(formatted_account)

        return eligible_account_batch

    def _process_batch(self, offset: int = 0) -> None:
        """Fetch the given batch of accounts and create worksheets for them."""
        eligible_accounts = self._get_eligible_account_batch_data(offset)
        PaymentBatchProcessor(
            self._abacus_event,
            eligible_accounts,
            append_vat_corrections=False,  # check payments should not have VAT corrections applied
        ).process()