"""Payment Batch Processor.""" from collections import defaultdict from typing import Dict, List, Optional from config import app_logger as logger from src.connectors.ows_payment import ( bulk_create_contract_payable_details, bulk_create_worksheet_contract_balance_after_tax, get_contract_balance_after_tax_entries, ) from src.constants import ( CLOSE_BALANCE_MSG, DECIMAL_ZERO, NO_ELIGIBLE_ACCOUNTS_ERR, POST_BALANCE_DETAILS_ENTRIES, POST_BALANCE_ENTRIES_AFTER_TAX, TaxCorrectionTypes, ) from src.models import ( Account, ContractCloseBalance, Event, PayableBalanceAfterTaxEntry, TaxCorrection, TaxCorrectionVAT, ) from src.processors.base.worksheet_calculator import WorksheetCalculator from src.processors.models import ( ContractLevelFlowthroughAllocationData, EligibleAccountLevelData, WorksheetPayableCalculatorDetail, ) from src.utils import ( fetch_all_contract_closing_balance_entries, fetch_all_flowthrough_allocation_entries, fetch_all_pending_tax_corrections, fetch_all_pending_tax_corrections_vat, ) class PaymentBatchProcessor: """Processor to calculate worksheets for all given accounts.""" def __init__( self, abacus_event: Event, accounts: List[EligibleAccountLevelData], payment_entity_policy_country_mapping: Optional[Dict[int, str]] = None, append_vat_corrections: Optional[bool] = True, ) -> None: """Init.""" self._abacus_event: Event = abacus_event self._statement_period_id = abacus_event.statement_period_id self._accounts_by_id: Dict[int, EligibleAccountLevelData] = { acc.account_id: acc for acc in accounts } self._contract_worksheets_by_id: Dict[int, WorksheetCalculator] = {} self._contract_statement_periods: Dict[int, int] = {} self._payment_entity_policy_country_mapping = ( payment_entity_policy_country_mapping or {} ) self._append_vat_corrections = append_vat_corrections def _get_account_ids(self) -> List[Account]: """Return list of account ids.""" ## TODO replace this with further refactoring, this is to get around current typing issues return [ Account(account_id=acc.account_id) for acc in self._accounts_by_id.values() ] def _get_contract_ids(self) -> List[int]: """Return list of contract ids.""" return list(self._contract_worksheets_by_id.keys()) def _get_closing_balance_related_data(self) -> List[ContractCloseBalance]: """Get closing balance data for account ids. Fetch missed items from previous statement period. """ accounts = self._get_account_ids() closing_balance_entries = fetch_all_contract_closing_balance_entries( self._statement_period_id, accounts ) closing_balance_account_ids = set( cb.account_id for cb in closing_balance_entries ) missed_accounts = [ account for account in accounts if account.account_id not in closing_balance_account_ids ] if missed_accounts: previous_closing_balances = fetch_all_contract_closing_balance_entries( self._statement_period_id - 1, missed_accounts ) closing_balance_entries.extend(previous_closing_balances) return closing_balance_entries def _setup_contract_worksheets(self) -> None: """Set up contract worksheet dictionary based on closed balance data.""" closing_balances = self._get_closing_balance_related_data() for balance in closing_balances: account = self._accounts_by_id[balance.account_id] # This is not prioritization, condition added to isolate from check processor where country_of_tax_policy is always US if ( not account.country_of_tax_policy and self._payment_entity_policy_country_mapping ): account.country_of_tax_policy = ( self._payment_entity_policy_country_mapping.get( balance.reference_payment_entity_id ) ) self._contract_worksheets_by_id[balance.contract_id] = WorksheetCalculator( balance, account ) self._contract_statement_periods = { balance.contract_id: balance.statement_period_id for balance in closing_balances } def _apply_wht_correction_to_worksheet(self, correction: TaxCorrection) -> None: """Apply a single WHT correction to its corresponding worksheet. Args: correction: The tax correction to apply """ wht_detail = WorksheetPayableCalculatorDetail( amount_payable=correction.amount, target_table='worksheet_tax_correction', target_id=correction.worksheet_tax_correction_id, payable_detail_type_id=correction.payable_detail_type_id, ) self._contract_worksheets_by_id[correction.contract_id].append_wht_item( wht_detail ) def _get_contracts_by_period(self) -> Dict[int, List[int]]: """Group contract IDs by their statement period. Returns: Dictionary mapping statement_period_id to list of contract_ids """ contracts_by_period: Dict[int, List[int]] = {} for contract_id in self._get_contract_ids(): statement_period_id = self._contract_statement_periods[contract_id] if statement_period_id not in contracts_by_period: contracts_by_period[statement_period_id] = [] contracts_by_period[statement_period_id].append(contract_id) return contracts_by_period def _append_wht_corrections_to_worksheets(self) -> None: """Fetch all pending wht corrections and associate with the current worksheets. Groups contracts by their statement period and fetches corrections per period. """ contracts_by_period = self._get_contracts_by_period() for statement_period_id, contract_ids in contracts_by_period.items(): wht_corrections = fetch_all_pending_tax_corrections( TaxCorrectionTypes.wht, contract_ids, statement_period_id, ) for correction in wht_corrections: self._apply_wht_correction_to_worksheet(correction) def _apply_vat_correction_to_worksheet(self, correction: TaxCorrectionVAT) -> None: """Apply a single VAT correction to its corresponding worksheet. Args: correction: The VAT tax correction to apply """ vat_detail = WorksheetPayableCalculatorDetail( amount_payable=correction.amount, target_table='worksheet_tax_correction_vat', target_id=correction.worksheet_tax_correction_vat_id, payable_detail_type_id=correction.payable_detail_type_id, ) self._contract_worksheets_by_id[correction.contract_id].append_vat_item( vat_detail ) def _append_vat_corrections_to_worksheets(self) -> None: """Fetch all pending vat corrections and associate with the current worksheets. Only processes VAT corrections when append_vat_corrections is True. """ if not self._append_vat_corrections: return contracts_by_period = self._get_contracts_by_period() for statement_period_id, contract_ids in contracts_by_period.items(): vat_corrections = fetch_all_pending_tax_corrections_vat( contract_ids, statement_period_id, ) for correction in vat_corrections: self._apply_vat_correction_to_worksheet(correction) def _get_created_worksheets(self) -> List[PayableBalanceAfterTaxEntry]: """Return list of created worksheet payable balance ids.""" contract_ids = self._get_contract_ids() batch_size = 100 # going for smaller batches for safety payable_after_tax_entries = [] start = 0 while start < len(contract_ids): end = start + batch_size contract_ids_for_query = contract_ids[start:end] response = get_contract_balance_after_tax_entries( self._abacus_event.abacus_event_id, limit=batch_size, contract_ids=contract_ids_for_query, ) payable_after_tax_entries.extend(response.items) start += batch_size return payable_after_tax_entries def _post_contract_worksheets_and_details(self) -> None: """Post contract worksheets and the associated details.""" formatted_worksheets = [] for worksheet in self._contract_worksheets_by_id.values(): formatted_worksheets.append(worksheet.get_payable_balance_after_tax_item()) if len(formatted_worksheets) < 1: return None logger.info(POST_BALANCE_ENTRIES_AFTER_TAX.format(len(formatted_worksheets))) bulk_create_worksheet_contract_balance_after_tax( self._abacus_event.abacus_event_id, self._statement_period_id, formatted_worksheets, ) payable_after_tax_entries = self._get_created_worksheets() formatted_details = [] for entry in payable_after_tax_entries: contract_id = entry.contract_id payable_after_tax_id = entry.worksheet_account_contract_payable_after_tax_id details = self._contract_worksheets_by_id[ contract_id ].get_payable_detail_items(payable_after_tax_id) formatted_details.extend(details) logger.info(POST_BALANCE_DETAILS_ENTRIES.format(len(formatted_details))) bulk_create_contract_payable_details( self._abacus_event.abacus_event_id, self._statement_period_id, formatted_details, ) def _append_flowthrough_allocations_to_worksheets(self) -> None: """Fetching flowthrough allocations entries and append to corresponding worksheets.""" allocations_by_contract = self._get_flowthrough_allocation_by_contract() for contract_id, contract_allocations in allocations_by_contract.items(): self._contract_worksheets_by_id[contract_id].apply_flowthrough_items( contract_allocations ) def _get_flowthrough_allocation_by_contract( self, ) -> Dict[int, ContractLevelFlowthroughAllocationData]: """Get mapped flowthrough allocation data for given contract ids.""" flowthrough_allocations = fetch_all_flowthrough_allocation_entries( self._get_contract_ids() ) formatted_allocations: Dict[int, ContractLevelFlowthroughAllocationData] = ( defaultdict( lambda: ContractLevelFlowthroughAllocationData( sum=DECIMAL_ZERO, items=[] ) ) ) for allocation in flowthrough_allocations: formatted_allocations[ allocation.contract_id ].sum += allocation.amount_to_payment formatted_allocations[allocation.contract_id].items.append(allocation) return formatted_allocations def process(self) -> None: """Calculate worksheets for the accounts in the class""" logger.info(CLOSE_BALANCE_MSG.format(len(self._accounts_by_id))) # no eligible accounts to process if not self._get_account_ids(): logger.info(NO_ELIGIBLE_ACCOUNTS_ERR.format(self._abacus_event.target_id)) return # generate initial worksheets based on just closing balance self._setup_contract_worksheets() # none of the eligible accounts had contract data if not self._get_contract_ids(): return self._append_flowthrough_allocations_to_worksheets() self._append_wht_corrections_to_worksheets() self._append_vat_corrections_to_worksheets() self._post_contract_worksheets_and_details()