Trixy7/trixy_core/satellite/loadbalancing/round_robin.py

# -*- coding: utf-8 -*-
"""
Round-Robin Load-Balancing-Strategie.

Verteilt Anfragen gleichmäßig auf alle verfügbaren Satellites.
"""

from typing import Any, TYPE_CHECKING

from trixy_core.satellite.loadbalancing.strategy import (
    LoadBalancingStrategy,
    SelectionResult,
)

if TYPE_CHECKING:
    from trixy_core.satellite.satellite import Satellite


class RoundRobinStrategy(LoadBalancingStrategy):
    """
    Round-Robin-Strategie.

    Wählt Satellites der Reihe nach aus und beginnt dann von vorne.
    Gewährleistet gleichmäßige Verteilung über alle Satellites.

    Example:
        strategy = RoundRobinStrategy()

        # Bei 3 Satellites: A, B, C
        # Auswahl-Reihenfolge: A, B, C, A, B, C, ...

        balancer = LoadBalancer(strategy)
        result = balancer.select(satellites)
    """

    def __init__(self) -> None:
        """Initialisiert die Round-Robin-Strategie."""
        self._current_index: int = 0
        self._satellite_order: list[str] = []

    @property
    def name(self) -> str:
        """Name der Strategie."""
        return "round_robin"

    def select(
        self,
        satellites: list["Satellite"],
        context: dict[str, Any] | None = None,
    ) -> SelectionResult | None:
        """
        Wählt den nächsten Satellite in der Rotation.

        Args:
            satellites: Liste verfügbarer Satellites
            context: Optionaler Kontext (wird ignoriert)

        Returns:
            SelectionResult oder None wenn keine Satellites verfügbar
        """
        if not satellites:
            return None

        # Sortiere Satellites für konsistente Reihenfolge
        sorted_satellites = sorted(satellites, key=lambda s: s.id)
        satellite_ids = [s.id for s in sorted_satellites]

        # Prüfe ob sich die Liste geändert hat
        if satellite_ids != self._satellite_order:
            self._satellite_order = satellite_ids
            # Index anpassen falls nötig
            if self._current_index >= len(sorted_satellites):
                self._current_index = 0

        # Satellite auswählen
        selected = sorted_satellites[self._current_index]

        # Index für nächste Auswahl vorbereiten
        self._current_index = (self._current_index + 1) % len(sorted_satellites)

        return SelectionResult(
            satellite=selected,
            reason=f"Round-Robin Index {self._current_index}",
            score=1.0 / len(sorted_satellites),
        )

    def reset(self) -> None:
        """Setzt den Index zurück."""
        self._current_index = 0
        self._satellite_order = []


class WeightedRoundRobinStrategy(LoadBalancingStrategy):
    """
    Gewichtetes Round-Robin.

    Berücksichtigt Gewichte bei der Rotation. Satellites mit
    höherem Gewicht werden häufiger ausgewählt.

    Example:
        strategy = WeightedRoundRobinStrategy()
        strategy.set_weight("satellite_1", 3)  # 3x so oft
        strategy.set_weight("satellite_2", 1)  # Normal
        strategy.set_weight("satellite_3", 2)  # 2x so oft

        # Reihenfolge: 1, 1, 1, 3, 3, 2, 1, 1, 1, ...
    """

    def __init__(self, default_weight: int = 1) -> None:
        """
        Initialisiert die gewichtete Round-Robin-Strategie.

        Args:
            default_weight: Standard-Gewicht für neue Satellites
        """
        self._default_weight = default_weight
        self._weights: dict[str, int] = {}
        self._current_index: int = 0
        self._expanded_list: list[str] = []

    @property
    def name(self) -> str:
        """Name der Strategie."""
        return "weighted_round_robin"

    def set_weight(self, satellite_id: str, weight: int) -> None:
        """
        Setzt das Gewicht für einen Satellite.

        Args:
            satellite_id: Satellite-ID
            weight: Gewicht (muss >= 1 sein)
        """
        self._weights[satellite_id] = max(1, weight)
        self._expanded_list = []  # Neuberechnung erzwingen

    def get_weight(self, satellite_id: str) -> int:
        """Gibt das Gewicht eines Satellites zurück."""
        return self._weights.get(satellite_id, self._default_weight)

    def _build_expanded_list(self, satellites: list["Satellite"]) -> list[str]:
        """Baut die expandierte Liste basierend auf Gewichten."""
        expanded = []
        for satellite in sorted(satellites, key=lambda s: s.id):
            weight = self.get_weight(satellite.id)
            expanded.extend([satellite.id] * weight)
        return expanded

    def select(
        self,
        satellites: list["Satellite"],
        context: dict[str, Any] | None = None,
    ) -> SelectionResult | None:
        """
        Wählt den nächsten Satellite basierend auf Gewichten.

        Args:
            satellites: Liste verfügbarer Satellites
            context: Optionaler Kontext

        Returns:
            SelectionResult oder None
        """
        if not satellites:
            return None

        # ID -> Satellite Mapping
        satellite_map = {s.id: s for s in satellites}

        # Expandierte Liste erstellen wenn nötig
        current_ids = sorted(s.id for s in satellites)
        if not self._expanded_list or set(self._expanded_list) != set(current_ids):
            self._expanded_list = self._build_expanded_list(satellites)
            if self._current_index >= len(self._expanded_list):
                self._current_index = 0

        if not self._expanded_list:
            return None

        # Auswahl
        selected_id = self._expanded_list[self._current_index]
        selected = satellite_map.get(selected_id)

        if selected is None:
            # Satellite nicht mehr verfügbar, Liste neu aufbauen
            self._expanded_list = self._build_expanded_list(satellites)
            self._current_index = 0
            if not self._expanded_list:
                return None
            selected_id = self._expanded_list[0]
            selected = satellite_map.get(selected_id)

        if selected is None:
            return None

        # Index vorrücken
        self._current_index = (self._current_index + 1) % len(self._expanded_list)

        return SelectionResult(
            satellite=selected,
            reason=f"Weighted Round-Robin (weight={self.get_weight(selected.id)})",
            score=self.get_weight(selected.id) / sum(self.get_weight(s.id) for s in satellites),
        )

    def reset(self) -> None:
        """Setzt den Zustand zurück."""
        self._current_index = 0
        self._expanded_list = []