from typing import Dict, Iterable, List, Optional

from src import constants
from src.connectors.ows_account import (
    account_payment_term_dataloader,
    get_account_tax_info,
    get_accounts,
)
from src.connectors.ows_payment import (
    get_contract_balance_after_tax_entries,
    get_contract_closing_balance_entries,
    get_contract_closing_balance_entries_bulk,
    get_flowthrough_allocation_entries,
    get_payable_details_entries,
    get_tax_corrections,
    get_tax_corrections_vat,
)
from src.connectors.utils import fetch_all
from src.models import (
    Account,
    AccountTaxInfo,
    ContractCloseBalance,
    Event,
    FlowthroughAllocation,
    PayableBalanceAfterTaxEntry,
    PayableDetailEntry,
    PaymentEntity,
    TaxCorrection,
    TaxCorrectionVAT,
)


def fetch_all_contract_closing_balance_entries(
    statement_period_id: int, accounts: List[Account]
) -> List[ContractCloseBalance]:
    """Get all contract closing balance entries for specified period."""
    offset = 0
    limit = constants.BATCH_SIZE
    closing_balance_entries: List[ContractCloseBalance] = []

    while True:
        batch = get_contract_closing_balance_entries(
            statement_period_id, accounts, limit=limit, offset=offset
        )

        offset = offset + limit
        closing_balance_entries.extend(batch.items)

        if offset >= batch.total_count:
            break

    return closing_balance_entries


def fetch_all_pending_tax_corrections(
    correction_type: constants.TaxCorrectionTypes,
    contract_ids: List[int],
    statement_period_id: Optional[int] = None,
) -> List[TaxCorrection]:
    """Fetch all pending tax corrections for contracts."""
    offset = 0
    limit = constants.BATCH_SIZE
    corrections: List[TaxCorrection] = []

    while True:
        batch = get_tax_corrections(
            correction_type,
            constants.TaxCorrectionStatuses.pending,
            statement_period_id,
            contract_ids,
            limit,
            offset,
        )

        offset += limit
        corrections.extend(batch.items)

        if offset >= batch.total_count:
            break

    return corrections


def fetch_all_pending_tax_corrections_vat(
    contract_ids: List[int], statement_period_id: Optional[int] = None
) -> List[TaxCorrectionVAT]:
    """Fetch all pending tax corrections VAT for contracts."""
    offset = 0
    limit = constants.BATCH_SIZE
    corrections: List[TaxCorrectionVAT] = []

    while True:
        batch = get_tax_corrections_vat(
            constants.TaxCorrectionStatuses.pending,
            statement_period_id,
            contract_ids,
            limit,
            offset,
        )

        offset += limit
        corrections.extend(batch.items)

        if offset >= batch.total_count:
            break

    return corrections


def fetch_all_contract_closing_balance_entries_bulk(
    worksheet_closing_balance_ids: List[int],
) -> List[ContractCloseBalance]:
    """Fetch all closing balance entries by closing balance ids."""
    return list(
        fetch_all(
            get_contract_closing_balance_entries_bulk,
            worksheet_closing_balance_ids,
            limit=constants.BATCH_SIZE_REFRESH,
        )
    )


def fetch_all_payable_details_entries(
    statement_period_id: int,
    worksheet_payable_after_tax_ids: List[int],
    detail_groups: List[str],
) -> List[PayableDetailEntry]:
    """Fetch all worksheet payable details for given after-tax ids and detail groups."""
    return list(
        fetch_all(
            get_payable_details_entries,
            statement_period_id,
            worksheet_payable_after_tax_ids,
            detail_groups,
        )
    )


def fetch_all_account_tax_info_entries(
    account_ids: Iterable[int],
) -> List[AccountTaxInfo]:
    """Get all account tax info entries."""
    offset = 0
    limit = constants.BATCH_SIZE
    entries: List[AccountTaxInfo] = []

    while True:
        batch = get_account_tax_info(list(account_ids), limit, offset)

        offset = offset + limit
        entries.extend(batch.items)

        if offset >= batch.total_count:
            break

    return entries


def create_payment_entity_to_country_of_tax_reporting_mapping(
    ref_payment_entities: List[PaymentEntity],
) -> Dict[int, str]:
    """Get country of tax policy for specified payment entity id."""
    mapping = dict()

    for entity in ref_payment_entities:
        payment_entity_id = entity.reference_payment_entity_id
        mapping[payment_entity_id] = entity.country_of_tax_reporting

    return mapping


def get_all_payable_balance_after_tax_entries_by_event(
    event: Event,
) -> List[PayableBalanceAfterTaxEntry]:
    """Get all payable balance after tax entries by abacus event."""
    offset = 0
    limit = constants.BATCH_SIZE_REFRESH
    entries: List[PayableBalanceAfterTaxEntry] = []

    while True:
        batch = get_contract_balance_after_tax_entries(
            event.abacus_event_id, limit=limit, offset=offset
        )

        offset = offset + limit
        entries.extend(batch.items)

        if offset >= batch.total_count:
            break

    return entries


def fetch_all_flowthrough_allocation_entries(
    contract_ids: List[int],
) -> List[FlowthroughAllocation]:
    """Get all flowthrough allocation entries."""
    offset = 0
    limit = constants.BATCH_SIZE
    allocations: List[FlowthroughAllocation] = []
    contract_ids_set = set(contract_ids)  # new variable here to pass type check

    while True:
        batch = get_flowthrough_allocation_entries(
            contract_ids=contract_ids_set, limit=limit, offset=offset
        )

        offset = offset + limit
        allocations.extend(batch.items)

        if offset >= batch.total_count:
            break

    return allocations