spirekeeper/calculations.py

"""
Pure calculation functions for DCA transaction processing.

These functions have no external dependencies (no lnbits, no database)
and can be easily tested in isolation.

What's intentionally NOT here (deleted 2026-05-26):
  - `calculate_commission` (back-derive principal+fee from a gross-with-
    commission-baked-in wire amount). Lamassu-era reverse-derivation;
    obsolete since bitSpire stamps `principal_sats` AND `fee_sats`
    directly on Payment.extra per aiolabs/lamassu-next#44.
  - `calculate_exchange_rate` (principal / fiat_amount). bitSpire stamps
    `exchange_rate` directly on Payment.extra too. Not used in production.
"""

from typing import Dict, Tuple


# Per-direction fee cap (super + operator) for any single direction.
# Locked at 15% per coord-log §2026-06-01T07:22Z (bitspire) — defense in
# depth: producer (this side) refuses to publish/persist > cap; consumer
# (bitspire) refuses to apply > cap. See aiolabs/satmachineadmin#37,#38
# and aiolabs/lamassu-next#57.
MAX_FEE_FRACTION_PER_DIRECTION = 0.15


def calculate_distribution(
    base_amount_sats: int,
    client_balances: Dict[str, float],
    min_balance_threshold: float = 0.01
) -> Dict[str, int]:
    """
    Calculate proportional distribution of sats to clients based on their fiat balances.

    Uses proportional allocation with remainder distribution to ensure
    the total distributed equals exactly the base amount.

    Args:
        base_amount_sats: Total sats to distribute (after commission)
        client_balances: Dict of {client_id: remaining_balance_fiat}
        min_balance_threshold: Minimum balance to be included (default 0.01)

    Returns:
        Dict of {client_id: allocated_sats}

    Example:
        >>> calculate_distribution(100000, {"a": 500.0, "b": 500.0})
        {"a": 50000, "b": 50000}
    """
    # Filter out clients with balance below threshold
    active_balances = {
        client_id: balance
        for client_id, balance in client_balances.items()
        if balance >= min_balance_threshold
    }

    if not active_balances:
        return {}

    total_balance = sum(active_balances.values())

    if total_balance == 0:
        return {}

    # First pass: calculate base allocations and track for remainder distribution
    client_calculations = []
    distributed_sats = 0

    for client_id, balance in active_balances.items():
        proportion = balance / total_balance
        exact_share = base_amount_sats * proportion
        allocated_sats = round(exact_share)

        client_calculations.append({
            'client_id': client_id,
            'proportion': proportion,
            'exact_share': exact_share,
            'allocated_sats': allocated_sats,
        })
        distributed_sats += allocated_sats

    # Handle remainder due to rounding
    remainder = base_amount_sats - distributed_sats

    if remainder != 0:
        # Sort by largest fractional remainder to distribute fairly
        client_calculations.sort(
            key=lambda x: x['exact_share'] - x['allocated_sats'],
            reverse=True
        )

        # Distribute remainder one sat at a time
        for i in range(abs(remainder)):
            idx = i % len(client_calculations)
            if remainder > 0:
                client_calculations[idx]['allocated_sats'] += 1
            else:
                client_calculations[idx]['allocated_sats'] -= 1

    # Build final distributions dict
    distributions = {
        calc['client_id']: calc['allocated_sats']
        for calc in client_calculations
    }

    return distributions


def split_principal_based(
    principal_sats: int,
    super_frac: float,
    operator_frac: float,
) -> Tuple[int, int]:
    """Compute platform + operator fee shares as independent fractions of
    `principal_sats`. Both shares are derived from the customer's
    principal (the canonical source of truth), NOT back-derived from
    `fee_sats`.

    Returns (platform_fee_sats, operator_fee_sats). Both are rounded
    independently; rounding remainders do NOT compound — the customer
    pays whatever bitspire collected, and any drift between (super +
    operator) and the bitspire-reported `fee_sats` surfaces via
    `dca_settlements.fee_mismatch_sats`.

    Examples:
        >>> split_principal_based(100_000, 0.03, 0.05)
        (3000, 5000)
        >>> split_principal_based(266_800, 0.03, 0.0)
        (8004, 0)
        >>> split_principal_based(100_000, 0.0, 0.0)
        (0, 0)
        >>> split_principal_based(100_000, 0.15, 0.0)
        (15000, 0)

    The pre-#38 bug this corrects: the old math interpreted the super
    fee as `fraction_of_fee` rather than `fraction_of_principal`. On a
    100_000-sat principal with an 8% total bitspire fee (= 8_000 sats
    fee_sats) and super_fraction=0.03, the bug paid the super
    `round(8_000 * 0.03) = 240` sats — ~13× below the intended
    `100_000 * 0.03 = 3_000` sats per-settlement. Repeated on every
    cash-out since the bitspire wire-shape landed. See
    aiolabs/satmachineadmin#37 (parent) + #38 (this layer).
    """
    if not (0.0 <= super_frac <= 1.0):
        raise ValueError(f"super_frac must be in [0, 1], got {super_frac}")
    if not (0.0 <= operator_frac <= 1.0):
        raise ValueError(f"operator_frac must be in [0, 1], got {operator_frac}")
    if principal_sats <= 0:
        return 0, 0
    platform = max(0, round(principal_sats * super_frac))
    operator = max(0, round(principal_sats * operator_frac))
    return platform, operator


def allocate_operator_split_legs(
    operator_fee_sats: int, leg_fractions: list
) -> list:
    """Stage-2 of the v2 commission split: the operator's remainder is sliced
    across N leg wallets per `leg_fractions` (each in [0, 1], sum should
    equal 1.0).

    The last leg absorbs the rounding remainder so the sum of allocations
    exactly equals operator_fee_sats (assuming fractions sum to ~1.0).
    Returns a list of integer sat amounts in the same order as leg_fractions.

    Examples:
        >>> allocate_operator_split_legs(70, [0.5, 0.3, 0.2])
        [35, 21, 14]
        >>> allocate_operator_split_legs(5575, [0.5, 0.3, 0.2])
        [2787, 1672, 1116]
        >>> allocate_operator_split_legs(100, [1.0])
        [100]
        >>> allocate_operator_split_legs(0, [0.5, 0.5])
        [0, 0]
    """
    if not leg_fractions:
        return []
    if operator_fee_sats <= 0:
        return [0] * len(leg_fractions)
    for f in leg_fractions:
        if not (0.0 <= float(f) <= 1.0):
            raise ValueError(
                f"every leg fraction must be in [0, 1], got {f}"
            )
    allocations: list = []
    remaining = operator_fee_sats
    for idx, fraction in enumerate(leg_fractions):
        if idx == len(leg_fractions) - 1:
            allocations.append(remaining)
        else:
            amount = round(operator_fee_sats * float(fraction))
            allocations.append(amount)
            remaining -= amount
    return allocations