"""Worksheet Calculator.""" from decimal import Decimal from typing import List from src.constants import DECIMAL_ZERO, PayableDetailTypes, TableNames from src.models import ( ContractCloseBalance, ContractPayableDetails, FlowthroughAllocation, PayableBalanceAfterTax, PayableBalanceAfterTaxEntry, ) from src.processors.models import ( ContractLevelFlowthroughAllocationData, EligibleAccountLevelData, WorksheetPayableCalculatorDetail, ) class WorksheetCalculator: """Represents a worksheet payment for a contract and associated payable details.""" def __init__( self, closed_balance: ContractCloseBalance, account_data: EligibleAccountLevelData, ) -> None: """Initialize the worksheet with the closing balance entry amount.""" self.closing_balance_id: int = ( closed_balance.worksheet_account_contract_closing_balance_id ) self.contract_id: int = closed_balance.contract_id self.account_id: int = closed_balance.account_id self.currency_code: str = closed_balance.currency_code self.payable_amount_pre_tax: Decimal = ( Decimal(closed_balance.amount) if closed_balance.amount > 0 else DECIMAL_ZERO ) self.tax_withholding_amount: Decimal | None = None self.vat_amount: Decimal | None = None self.payable_amount_post_tax: Decimal = self.payable_amount_pre_tax self.country_of_tax_residence: str = account_data.country_of_tax_residence self.country_of_tax_policy: str | None = account_data.country_of_tax_policy # make this a class attribut for ft to eaiser override in self._details passing value by address self._closing_balance_detail = WorksheetPayableCalculatorDetail( amount_payable=self.payable_amount_pre_tax, target_table='worksheet_account_contract_closing_balance', target_id=self.closing_balance_id, payable_detail_type_id=PayableDetailTypes.closing_balance, ) self._details: List[WorksheetPayableCalculatorDetail] = [ self._closing_balance_detail ] @classmethod def from_payable_balance_after_tax_entry( cls, entry: PayableBalanceAfterTaxEntry ) -> 'WorksheetCalculator': """Build a calculator from an existing payable balance after tax entry. Used by the refresh flow: seeds state from the persisted entry and starts with an empty details list so that only the rebuilt correction details (wht/vat) are emitted by get_payable_detail_items. Mirrors the attribute setup in __init__. """ instance = cls.__new__(cls) instance.closing_balance_id = ( entry.worksheet_account_contract_closing_balance_id ) instance.contract_id = entry.contract_id instance.account_id = entry.account_id instance.currency_code = entry.currency_code instance.payable_amount_pre_tax = entry.payable_amount_pre_tax instance.tax_withholding_amount = None instance.vat_amount = None instance.payable_amount_post_tax = entry.payable_amount_pre_tax instance.country_of_tax_residence = entry.country_of_tax_residence instance.country_of_tax_policy = entry.country_of_tax_policy instance._details = [] return instance def _recalculate_payable_amount_post_tax(self) -> None: """Recalculate post tax payable amount.""" self.payable_amount_post_tax = self.payable_amount_pre_tax self.payable_amount_post_tax += self.tax_withholding_amount or DECIMAL_ZERO self.payable_amount_post_tax += self.vat_amount or DECIMAL_ZERO def append_wht_item(self, wht_detail: WorksheetPayableCalculatorDetail) -> None: """Append new wht item to worksheet and recalculate values.""" if self.tax_withholding_amount is None: self.tax_withholding_amount = DECIMAL_ZERO self.tax_withholding_amount += Decimal(wht_detail.amount_payable) self._details.append(wht_detail) self._recalculate_payable_amount_post_tax() def append_vat_item(self, vat_detail: WorksheetPayableCalculatorDetail) -> None: """Append new vat item to worksheet and recalculate values.""" if self.vat_amount is None: self.vat_amount = DECIMAL_ZERO self.vat_amount += Decimal(vat_detail.amount_payable) self._details.append(vat_detail) self._recalculate_payable_amount_post_tax() def apply_flowthrough_items( self, contract_allocations: ContractLevelFlowthroughAllocationData ) -> None: """Apply new flowthrough allocation items to worksheet and recalculate values.""" ft_sum, ft_items = contract_allocations.sum, contract_allocations.items # Only skip when there is no flowthrough allocation items. if not ft_items: return if self.payable_amount_pre_tax >= ft_sum: self._append_flowthrough_items(ft_items, set_payable_zero=True) else: self._override_payable_amounts( ft_sum ) # This step overrides payable_amount_pre_tax and payable_amount_post_tax to flowthrough sum. self._override_payable_details_closing_balance_entry( DECIMAL_ZERO ) # This step overrides closing balance payable details amount_payable to 0. self._append_flowthrough_items(ft_items) def _append_flowthrough_items( self, items: List[FlowthroughAllocation], set_payable_zero: bool = False ) -> None: """Append flowthrough allocation items to worksheet details list, amount_payable will be set to 0 if set_payable_zero is True.""" for item in items: ft_detail = WorksheetPayableCalculatorDetail( amount_payable=DECIMAL_ZERO if set_payable_zero else item.amount_to_payment, target_table=TableNames.payment_allocation, target_id=item.payment_allocation_id, payable_detail_type_id=PayableDetailTypes.flowthrough_allocation, ) self._details.append(ft_detail) def _override_payable_amounts(self, amount: Decimal) -> None: """Override payable_amount_pre_tax and payable_amount_post_tax to the given amount.""" self.payable_amount_pre_tax = amount self._recalculate_payable_amount_post_tax() def _override_payable_details_closing_balance_entry(self, amount: Decimal) -> None: """Override closing balance payable details amount_payable to the given amount.""" self._closing_balance_detail.amount_payable = amount def get_payable_balance_after_tax_item(self) -> PayableBalanceAfterTax: """Return the item we will use for posting to ows-payment The payable balance after tax item seems incredibly overloaded. So rather making it an integral part of this class, I decided to put this here thinking we could refactor it later. """ return PayableBalanceAfterTax( worksheet_account_contract_closing_balance_id=self.closing_balance_id, contract_id=self.contract_id, account_id=self.account_id, currency_code=self.currency_code, payable_amount_pre_tax=self.payable_amount_pre_tax, tax_withholding_amount=self.tax_withholding_amount, vat_amount=self.vat_amount, payable_amount_post_tax=self.payable_amount_post_tax, country_of_tax_residence=self.country_of_tax_residence, country_of_tax_policy=self.country_of_tax_policy, ) def get_payable_detail_items( self, worksheet_payable_after_tax_id: int ) -> List[ContractPayableDetails]: """Get payable detail items.""" formatted_entries = [] for item in self._details: item_dict = item.model_dump() item_dict.update( { 'worksheet_account_contract_payable_after_tax_id': worksheet_payable_after_tax_id, 'account_id': self.account_id, 'contract_id': self.contract_id, 'currency': self.currency_code, } ) formatted_entries.append(ContractPayableDetails.model_validate(item_dict)) return formatted_entries