from datetime import datetime

from app.models import CampaignBatch
from app.strategies.base import QuotaAllocationStrategy
from app.strategies.types import QuotaAllocation, QuotaStrategyType


class FairDistributionStrategy(QuotaAllocationStrategy):
    """
    Fair distribution quota allocation strategy.

    Distributes available quota evenly among all active campaign batches,
    with minute-based rotation to ensure fairness over time.
    """

    strategy_type = QuotaStrategyType.FAIR_DISTRIBUTION

    def allocate_quota(
        self,
        batches: list[CampaignBatch],
        available_quota: int,
        current_time: datetime,
    ) -> list[QuotaAllocation]:
        """
        Allocate quota fairly and efficiently among all batches.

        Algorithm:
        1. Rotate batch order based on current minute for fairness over time.
        2. Use water-filling to distribute quota: smaller-capacity batches are
           capped first, remaining quota is split equally among the rest.
        """
        if not batches or available_quota <= 0:
            return []

        batches = self._rotate_batches(batches, current_time)
        allocations = self._distribute_quota(batches, available_quota)

        return [
            QuotaAllocation(
                batch_id=batch.id,
                campaign_id=batch.campaign_id,
                emails_to_send=emails_to_send,
            )
            for batch in batches
            if (emails_to_send := allocations.get(batch.id, 0)) > 0
        ]

    @staticmethod
    def _distribute_quota(
        batches: list[CampaignBatch], available_quota: int
    ) -> dict[int, int]:
        """Distribute quota fairly using a single-pass water-filling approach.

        Batches are sorted by remaining capacity. Starting from the smallest,
        each batch is checked: if its capacity fits within an equal share of the
        remaining quota, it is capped at its full capacity. Otherwise, all
        remaining batches receive an equal base allocation, with extra units
        distributed in rotation order for fairness.
        """
        capacities = {
            batch.id: max(0, batch.batch_size - batch.batch_offset) for batch in batches
        }
        sorted_by_capacity = sorted(batches, key=lambda b: capacities[b.id])

        allocations: dict[int, int] = {}
        remaining = available_quota
        capped_batch_ids: set[int] = set()

        for i, batch in enumerate(sorted_by_capacity):
            batches_left = len(sorted_by_capacity) - i
            fair_share = remaining // batches_left

            if capacities[batch.id] <= fair_share:
                allocations[batch.id] = capacities[batch.id]
                remaining -= capacities[batch.id]
                capped_batch_ids.add(batch.id)
            else:
                base = remaining // batches_left
                extra = remaining % batches_left
                for uncapped in sorted_by_capacity[i:]:
                    allocations[uncapped.id] = base
                for rotated in batches:
                    if extra <= 0:
                        break
                    if rotated.id not in capped_batch_ids:
                        allocations[rotated.id] += 1
                        extra -= 1
                return allocations

        return allocations

    @staticmethod
    def _rotate_batches(
        batches: list[CampaignBatch], current_time: datetime
    ) -> list[CampaignBatch]:
        """Rotate batch order based on current minute to ensure fairness over time."""
        if len(batches) <= 1:
            return batches

        # Sort by update time to ensure fair ordering
        batches = sorted(batches, key=lambda batch: batch.updated_at)

        # Rotate by current minute
        current_minute = current_time.minute
        remainder = current_minute % len(batches)
        return batches[remainder:] + batches[:remainder]