"""Payment Allocation Processor."""

from __future__ import annotations

from collections import defaultdict
from decimal import Decimal
from typing import Any

from lambdacommon.common_config import logger

from config import config
from src.connectors import ows_royalties
from src.connectors.repository import Repository
from src.enums import (
    CloseBalanceActionName,
    CloseBalanceActionStatus,
    PayeeType,
    PaymentAllocationType,
)
from src.errors import (
    BalancesNotClosedError,
    OwsServiceException,
    StatementPeriodPaymentEntityNotFoundError,
)
from src.schemas import (
    LedgerAdjustmentApplied,
    PaymentAllocationResponse,
    StatementPeriodPaymentEntity,
)
from src.utils import DEFAULT_SAFETY_MARGIN_PCT, build_batches, calculate_max_batch_size

# Grouping key: (contract_id, account_payee_id, adjustment_currency_code, payee_currency_code)
GroupKey = tuple[int, int, str, str]

# Worst-case query size estimation for max_allowed_packet clamping
_BASE_QUERY_BYTES = 1000  # Conservative estimate of query template size
_BYTES_PER_CONTRACT_ID = 20  # Max digits for a BIGINT value


class PaymentAllocationProcessor:
    """Payment allocation processor.

    Queries ledger_adjustment_applied for flowthrough adjustments in batches,
    groups them by contract/payee/currency, and creates payment_allocation
    records with linked payment_allocation_ledger_adjustment records.
    Commits after each batch for resilience.
    """

    def __init__(self, repository: Repository) -> None:
        """Initialize processor.

        Args:
            repository: Repository instance
        """
        self._repository = repository

    def process(self, sppe_id: int) -> PaymentAllocationResponse:
        """Process payment allocations for a statement period payment entity.

        Validates preconditions, then processes adjustments in contract-based
        batches. The NOT EXISTS clause in the queries ensures already-linked
        adjustments are skipped, making re-runs safe.

        Args:
            sppe_id: The statement_period_payment_entity_id to process

        Returns:
            PaymentAllocationResponse with counts of records created

        Raises:
            StatementPeriodPaymentEntityNotFoundError: If SPPE not found
            BalancesNotClosedError: If close_balance is not complete
            TransientError: If database connection fails
        """
        # 1. Look up SPPE
        logger.info('Looking up statement period payment entity')
        sppe = self._get_statement_period_payment_entity(sppe_id)
        logger.info(
            f'Found SPPE {sppe_id} with '
            f'statement_period_id={sppe.statement_period_id}, '
            f'reference_payment_entity_id={sppe.reference_payment_entity_id}'
        )

        # 2. Validate preconditions
        logger.info('Validating close_balance state')
        self._validate_balances_closed(sppe_id)

        # 3. Process adjustments
        allocations_created, total_linked = self._process_adjustments(
            sppe.statement_period_id, sppe.reference_payment_entity_id
        )

        return PaymentAllocationResponse(
            statement_period_id=sppe.statement_period_id,
            statement_period_payment_entity_id=sppe_id,
            allocations_created=allocations_created,
            ledger_adjustments_linked=total_linked,
        )

    def _create_allocations(
        self,
        grouped: dict[GroupKey, list[LedgerAdjustmentApplied]],
        statement_period_id: int,
    ) -> tuple[int, int]:
        """Create payment_allocation and link records for grouped adjustments.

        Args:
            grouped: Adjustments grouped by (contract_id, payee_id, currency)
            statement_period_id: The statement period ID

        Returns:
            Tuple of (allocations_created, total_linked)
        """
        allocations_created = 0
        total_linked = 0

        for group_key, group_adjustments in grouped.items():
            contract_id, account_payee_id, currency_code, _payee_currency_code = (
                group_key
            )

            amount_to_payment = sum(
                (adj.adjustment_amount_payee_currency for adj in group_adjustments),
                Decimal(0),
            )
            amount_to_ledger = sum(
                (adj.adjustment_amount for adj in group_adjustments),
                Decimal(0),
            )

            pa_id = self._repository.create_payment_allocation(
                contract_id=contract_id,
                payee_type=PayeeType.ACCOUNT_PAYEE,
                payee_id=account_payee_id,
                statement_period_id=statement_period_id,
                payment_allocation_type=PaymentAllocationType.FLOWTHROUGH,
                amount_to_payment=amount_to_payment,
                amount_to_ledger=amount_to_ledger,
                currency_code=currency_code,
                description=config.payment_allocation_description,
                created_by=config.app_name,
            )

            for adj in group_adjustments:
                self._repository.create_payment_allocation_ledger_adjustment(
                    pa_id, adj.ledger_adjustment_applied_id, config.app_name
                )
                total_linked += 1

            allocations_created += 1

        return allocations_created, total_linked

    def _get_effective_batch_size(self) -> int:
        """Calculate effective batch_size, clamped by max_allowed_packet.

        Queries the MySQL server's max_allowed_packet and estimates the
        worst-case query size for get_adjustments_for_contracts assuming
        batch_size contract IDs in the IN clause. If the query would
        exceed the packet limit, batch_size is reduced.

        Returns:
            The effective batch_size to use for processing.
        """
        max_bytes = self._repository.get_max_allowed_packet()
        if max_bytes is None:
            return config.batch_size

        max_contracts = calculate_max_batch_size(
            max_query_bytes=max_bytes,
            bytes_per_entry=_BYTES_PER_CONTRACT_ID,
            base_query_bytes=_BASE_QUERY_BYTES,
            safety_margin_pct=DEFAULT_SAFETY_MARGIN_PCT,
        )

        if max_contracts == 0:
            raise ValueError(
                f'max_allowed_packet={max_bytes} is too small for '
                f'the adjustment query (base={_BASE_QUERY_BYTES} bytes)'
            )

        effective = min(config.batch_size, max_contracts)
        if effective < config.batch_size:
            logger.warning(
                f'batch_size clamped from {config.batch_size} to {effective} '
                f'due to max_allowed_packet={max_bytes}'
            )
        return effective

    def _get_statement_period_payment_entity(
        self, sppe_id: int
    ) -> StatementPeriodPaymentEntity:
        """Look up statement period payment entity from the database.

        Args:
            sppe_id: The statement period payment entity ID

        Returns:
            The StatementPeriodPaymentEntity record

        Raises:
            StatementPeriodPaymentEntityNotFoundError: If not found
        """
        sppe = self._repository.get_statement_period_payment_entity(sppe_id)

        if sppe is None:
            raise StatementPeriodPaymentEntityNotFoundError(
                f'Statement period payment entity not found: '
                f'statement_period_payment_entity_id={sppe_id}'
            )

        return sppe

    def _process_adjustment_batch(
        self,
        adjustments: list[LedgerAdjustmentApplied],
        statement_period_id: int,
    ) -> tuple[int, int]:
        """Group adjustments, create allocations, and commit.

        Args:
            adjustments: List of adjustment records to process
            statement_period_id: The statement period ID

        Returns:
            Tuple of (allocations_created, total_linked)
        """
        grouped = self._group_adjustments(adjustments)
        alloc, linked = self._create_allocations(grouped, statement_period_id)
        self._repository.conn.commit()
        return alloc, linked

    def _process_adjustments(
        self,
        statement_period_id: int,
        reference_payment_entity_id: int,
    ) -> tuple[int, int]:
        """Process adjustments in contract-based batches.

        Outer loop discovers contracts with unlinked adjustments. Inner loop
        bin-packs contracts into batches and processes each. All batches use
        LIMIT with a counter to avoid extra empty-result queries. Commits
        after each batch.

        Args:
            statement_period_id: The statement period ID
            reference_payment_entity_id: The reference payment entity ID

        Returns:
            Tuple of (allocations_created, total_linked)
        """
        batch_size = self._get_effective_batch_size()
        allocations_created = 0
        total_linked = 0
        batch_number = 0

        while True:
            contract_counts = self._repository.get_contract_adjustment_counts(
                statement_period_id,
                reference_payment_entity_id,
                config.contract_batch_size,
            )
            if not contract_counts:
                break

            logger.info(
                f'Discovered {len(contract_counts)} contracts with unlinked adjustments'
            )

            batches = build_batches(
                contract_counts,
                batch_size,
                key=lambda cc: cc.contract_id,
                size=lambda cc: cc.adjustment_count,
            )

            alloc, linked, batch_number = self._process_batches(
                batches,
                batch_size,
                statement_period_id,
                reference_payment_entity_id,
                batch_number,
            )
            allocations_created += alloc
            total_linked += linked

        logger.info(
            f'Finished: created {allocations_created} allocations, '
            f'linked {total_linked} ledger adjustments '
            f'in {batch_number} batch(es)'
        )

        return allocations_created, total_linked

    def _process_batches(
        self,
        batches: list[list[tuple[int, int]]],
        batch_size: int,
        statement_period_id: int,
        reference_payment_entity_id: int,
        batch_number: int,
    ) -> tuple[int, int, int]:
        """Process contract batches with LIMIT, querying until empty.

        Each batch is a list of (contract_id, adjustment_count) pairs.
        Queries repeatedly with LIMIT until no unlinked adjustments remain,
        guaranteeing all adjustments are processed regardless of count accuracy.

        Args:
            batches: List of (contract_id, count) lists from build_batches
            batch_size: Maximum adjustments per LIMIT query
            statement_period_id: The statement period ID
            reference_payment_entity_id: The reference payment entity ID
            batch_number: Running batch counter for logging

        Returns:
            Tuple of (allocations_created, total_linked, batch_number)
        """
        allocations_created = 0
        total_linked = 0

        for batch in batches:
            contract_ids = [cid for cid, _ in batch]

            while True:
                batch_number += 1
                adjustments = self._repository.get_adjustments_for_contracts(
                    statement_period_id,
                    reference_payment_entity_id,
                    contract_ids,
                    limit=batch_size,
                )
                if not adjustments:
                    break

                logger.info(
                    f'Batch {batch_number}: {len(adjustments)} adjustments '
                    f'across {len(contract_ids)} contract(s)'
                )

                alloc, linked = self._process_adjustment_batch(
                    adjustments, statement_period_id
                )
                allocations_created += alloc
                total_linked += linked
                logger.info(
                    f'Batch {batch_number}: created {alloc} allocations, '
                    f'linked {linked} ledger adjustments'
                )

        return allocations_created, total_linked, batch_number

    def _validate_balances_closed(self, sppe_id: int) -> None:
        """Validate that close_balance action is complete for the SPPE.

        Queries the abacus_state table directly via the repository.

        Args:
            sppe_id: The statement period payment entity ID

        Raises:
            BalancesNotClosedError: If close_balance is not complete
        """
        # Workaround: query DB directly to avoid OWS timeout
        status = self._repository.get_state_status(
            'statement_period_payment_entity',
            sppe_id,
            CloseBalanceActionName.CLOSE_BALANCE,
        )
        if status != CloseBalanceActionStatus.COMPLETE:
            raise BalancesNotClosedError(
                f'Close balance is not complete for '
                f'statement_period_payment_entity_id={sppe_id}. '
                f'Balances must be closed before payment allocation'
            )

        logger.info('Close balance state is complete')

    def _validate_balances_closed_ows(
        self, statement_period_id: int, sppe_id: int
    ) -> None:
        """Validate that close_balance action is complete for the SPPE via OWS.

        Args:
            statement_period_id: The statement period ID
            sppe_id: The statement period payment entity ID

        Raises:
            BalancesNotClosedError: If close_balance is not complete
            OwsServiceException: If OWS call fails
        """
        response = ows_royalties.get_statement_period_payment_entity_states(
            statement_period_id
        )

        if response.status_code != 200:
            raise OwsServiceException(
                f'Failed to get payment entity states for '
                f'statement_period_id={statement_period_id}. '
                f'Status: {response.status_code}'
            )

        states: list[dict[str, Any]] = response.json()

        close_balance_state = None
        for state in states:
            if (
                state.get('statement_period_payment_entity_id') == sppe_id
                and state.get('action_name') == CloseBalanceActionName.CLOSE_BALANCE
            ):
                close_balance_state = state
                break

        if (
            not close_balance_state
            or close_balance_state.get('action_status')
            != CloseBalanceActionStatus.COMPLETE
        ):
            raise BalancesNotClosedError(
                f'Close balance is not complete for '
                f'statement_period_payment_entity_id={sppe_id}. '
                f'Balances must be closed before payment allocation'
            )

        logger.info('Close balance state is complete')

    @staticmethod
    def _group_adjustments(
        adjustments: list[LedgerAdjustmentApplied],
    ) -> dict[GroupKey, list[LedgerAdjustmentApplied]]:
        """Group adjustments by (contract_id, account_payee_id, currency_code, payee_currency_code).

        Args:
            adjustments: List of adjustment records

        Returns:
            Dict mapping group keys to lists of adjustments
        """
        grouped: dict[GroupKey, list[LedgerAdjustmentApplied]] = defaultdict(list)

        for adj in adjustments:
            key: GroupKey = (
                adj.contract_id,
                adj.account_payee_id,
                adj.adjustment_currency_code,
                adj.adjustment_payee_currency_code,
            )
            grouped[key].append(adj)

        return dict(grouped)