Trixy7/trixy_core/satellite/loadbalancing/least_connections.py

# -*- coding: utf-8 -*-
"""
Least-Connections Load-Balancing-Strategie.

Wählt den Satellite mit den wenigsten aktiven Verbindungen/Anfragen.
"""

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 LeastConnectionsStrategy(LoadBalancingStrategy):
    """
    Least-Connections-Strategie.

    Wählt immer den Satellite mit den wenigsten aktiven Anfragen.
    Ideal für Anfragen mit unterschiedlicher Bearbeitungsdauer.

    Example:
        strategy = LeastConnectionsStrategy()
        balancer = LoadBalancer(strategy)

        result = balancer.select(satellites)
        satellite = result.satellite

        # Anfrage starten (erhöht Connection-Count)
        await satellite.speak("Hallo")

        # Anfrage beenden (verringert Connection-Count)
        balancer.end_request(satellite)
    """

    def __init__(self) -> None:
        """Initialisiert die Least-Connections-Strategie."""
        self._connections: dict[str, int] = {}

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

    def select(
        self,
        satellites: list["Satellite"],
        context: dict[str, Any] | None = None,
    ) -> SelectionResult | None:
        """
        Wählt den Satellite mit den wenigsten Verbindungen.

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

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

        # Sortiere nach Verbindungen (aufsteigend), dann nach ID für Konsistenz
        sorted_satellites = sorted(
            satellites,
            key=lambda s: (self.get_connection_count(s.id), s.id)
        )

        selected = sorted_satellites[0]
        connections = self.get_connection_count(selected.id)

        return SelectionResult(
            satellite=selected,
            reason=f"Least connections ({connections})",
            score=1.0 / max(1, connections + 1),
            metadata={"connection_count": connections}
        )

    def get_connection_count(self, satellite_id: str) -> int:
        """
        Gibt die Anzahl aktiver Verbindungen zurück.

        Args:
            satellite_id: Satellite-ID

        Returns:
            Anzahl Verbindungen
        """
        return self._connections.get(satellite_id, 0)

    def on_request_start(self, satellite: "Satellite") -> None:
        """Erhöht den Connection-Count."""
        self._connections[satellite.id] = self.get_connection_count(satellite.id) + 1

    def on_request_end(self, satellite: "Satellite", success: bool = True) -> None:
        """Verringert den Connection-Count."""
        current = self.get_connection_count(satellite.id)
        if current > 0:
            self._connections[satellite.id] = current - 1

    def reset(self) -> None:
        """Setzt alle Connection-Counts zurück."""
        self._connections.clear()

    def get_all_counts(self) -> dict[str, int]:
        """Gibt alle Connection-Counts zurück."""
        return self._connections.copy()


class LeastResponseTimeStrategy(LoadBalancingStrategy):
    """
    Least-Response-Time-Strategie.

    Kombiniert Verbindungsanzahl mit durchschnittlicher Antwortzeit.
    Wählt den Satellite mit der niedrigsten erwarteten Wartezeit.

    Formel: score = connections * avg_response_time

    Example:
        strategy = LeastResponseTimeStrategy()
        balancer = LoadBalancer(strategy)

        result = balancer.select(satellites)
        satellite = result.satellite

        start_time = time.monotonic()
        await satellite.speak("Hallo")
        response_time = time.monotonic() - start_time

        strategy.record_response_time(satellite.id, response_time)
        balancer.end_request(satellite)
    """

    def __init__(self, window_size: int = 10) -> None:
        """
        Initialisiert die Strategie.

        Args:
            window_size: Anzahl der letzten Antwortzeiten für Durchschnitt
        """
        self._connections: dict[str, int] = {}
        self._response_times: dict[str, list[float]] = {}
        self._window_size = window_size

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

    def record_response_time(self, satellite_id: str, response_time: float) -> None:
        """
        Zeichnet eine Antwortzeit auf.

        Args:
            satellite_id: Satellite-ID
            response_time: Antwortzeit in Sekunden
        """
        if satellite_id not in self._response_times:
            self._response_times[satellite_id] = []

        times = self._response_times[satellite_id]
        times.append(response_time)

        # Nur letzte N Werte behalten
        if len(times) > self._window_size:
            self._response_times[satellite_id] = times[-self._window_size:]

    def get_avg_response_time(self, satellite_id: str) -> float:
        """
        Gibt die durchschnittliche Antwortzeit zurück.

        Args:
            satellite_id: Satellite-ID

        Returns:
            Durchschnittliche Antwortzeit (Standard: 1.0)
        """
        times = self._response_times.get(satellite_id, [])
        if not times:
            return 1.0  # Standard-Annahme
        return sum(times) / len(times)

    def get_connection_count(self, satellite_id: str) -> int:
        """Gibt die Anzahl aktiver Verbindungen zurück."""
        return self._connections.get(satellite_id, 0)

    def _calculate_score(self, satellite_id: str) -> float:
        """Berechnet den Score für einen Satellite (niedriger = besser)."""
        connections = self.get_connection_count(satellite_id)
        avg_time = self.get_avg_response_time(satellite_id)
        # Formel: Aktuelle Last * durchschnittliche Zeit
        return (connections + 1) * avg_time

    def select(
        self,
        satellites: list["Satellite"],
        context: dict[str, Any] | None = None,
    ) -> SelectionResult | None:
        """
        Wählt den Satellite mit der niedrigsten erwarteten Wartezeit.

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

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

        # Sortiere nach Score (aufsteigend)
        sorted_satellites = sorted(
            satellites,
            key=lambda s: (self._calculate_score(s.id), s.id)
        )

        selected = sorted_satellites[0]
        score = self._calculate_score(selected.id)

        return SelectionResult(
            satellite=selected,
            reason=f"Least response time (score={score:.3f})",
            score=1.0 / max(0.001, score),
            metadata={
                "connection_count": self.get_connection_count(selected.id),
                "avg_response_time": self.get_avg_response_time(selected.id),
                "calculated_score": score,
            }
        )

    def on_request_start(self, satellite: "Satellite") -> None:
        """Erhöht den Connection-Count."""
        self._connections[satellite.id] = self.get_connection_count(satellite.id) + 1

    def on_request_end(self, satellite: "Satellite", success: bool = True) -> None:
        """Verringert den Connection-Count."""
        current = self.get_connection_count(satellite.id)
        if current > 0:
            self._connections[satellite.id] = current - 1

    def reset(self) -> None:
        """Setzt den Zustand zurück."""
        self._connections.clear()
        self._response_times.clear()