Trixy7/trixy_core/network/rpc/correlation.py

# -*- coding: utf-8 -*-
"""
Request-Response Korrelation.

Verwaltet die Zuordnung von Anfragen zu Antworten
über eindeutige Korrelations-IDs.
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum, auto
from typing import Any, Awaitable, Callable, Generic, TypeVar

T = TypeVar("T")


class RequestState(Enum):
    """Zustand einer ausstehenden Anfrage."""
    PENDING = auto()      # Wartet auf Antwort
    COMPLETED = auto()    # Antwort erhalten
    TIMEOUT = auto()      # Timeout
    CANCELLED = auto()    # Abgebrochen
    ERROR = auto()        # Fehler


@dataclass
class CorrelationConfig:
    """Konfiguration für Korrelations-Manager."""
    default_timeout: float = 30.0     # Standard-Timeout in Sekunden
    cleanup_interval: float = 60.0    # Interval für Cleanup alter Requests
    max_pending: int = 10000          # Maximale Anzahl ausstehender Requests
    id_generator: Callable[[], str] | None = None  # Custom ID-Generator


@dataclass
class PendingRequest(Generic[T]):
    """
    Repräsentiert eine ausstehende Anfrage.

    Enthält alle Informationen über den Request und
    ermöglicht das Warten auf die Antwort.
    """
    correlation_id: str
    created_at: datetime = field(default_factory=datetime.now)
    timeout: float = 30.0
    state: RequestState = RequestState.PENDING
    result: T | None = None
    error: Exception | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    _future: asyncio.Future | None = field(default=None, repr=False)

    @property
    def is_complete(self) -> bool:
        """Prüft ob Request abgeschlossen ist."""
        return self.state in (
            RequestState.COMPLETED,
            RequestState.TIMEOUT,
            RequestState.CANCELLED,
            RequestState.ERROR,
        )

    @property
    def elapsed(self) -> timedelta:
        """Vergangene Zeit seit Erstellung."""
        return datetime.now() - self.created_at

    @property
    def elapsed_seconds(self) -> float:
        """Vergangene Zeit in Sekunden."""
        return self.elapsed.total_seconds()

    def is_expired(self) -> bool:
        """Prüft ob Request abgelaufen ist."""
        return self.elapsed_seconds > self.timeout


class CorrelationManager(Generic[T]):
    """
    Manager für Request-Response Korrelation.

    Verwaltet ausstehende Anfragen und ordnet
    eingehende Antworten den richtigen Requests zu.

    Beispiel:
        manager = CorrelationManager[bytes]()

        # Request erstellen
        correlation_id = await manager.create_request()

        # Nachricht senden mit correlation_id
        await send_message(data, correlation_id)

        # Auf Antwort warten
        try:
            response = await manager.wait_for_response(correlation_id)
        except TimeoutError:
            print("Timeout")

        # Antwort empfangen (von anderem Task)
        manager.complete(correlation_id, response_data)
    """

    def __init__(
        self,
        config: CorrelationConfig | None = None,
        logger: logging.Logger | None = None,
    ) -> None:
        """
        Initialisiert den Korrelations-Manager.

        Args:
            config: Konfiguration
            logger: Logger-Instanz
        """
        self.config = config or CorrelationConfig()
        self.logger = logger or logging.getLogger(__name__)

        self._pending: dict[str, PendingRequest[T]] = {}
        self._lock = asyncio.Lock()
        self._cleanup_task: asyncio.Task | None = None

        # ID-Generator
        self._generate_id = (
            self.config.id_generator
            or self._default_id_generator
        )

    @staticmethod
    def _default_id_generator() -> str:
        """Standard-ID-Generator."""
        return str(uuid.uuid4())

    @property
    def pending_count(self) -> int:
        """Anzahl ausstehender Requests."""
        return len(self._pending)

    async def start_cleanup(self) -> None:
        """Startet den Cleanup-Task."""
        if self._cleanup_task is None:
            self._cleanup_task = asyncio.create_task(self._cleanup_loop())

    async def stop_cleanup(self) -> None:
        """Stoppt den Cleanup-Task."""
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
            self._cleanup_task = None

    async def _cleanup_loop(self) -> None:
        """Periodisches Cleanup abgelaufener Requests."""
        while True:
            try:
                await asyncio.sleep(self.config.cleanup_interval)
                await self.cleanup_expired()
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"Cleanup-Fehler: {e}")

    async def create_request(
        self,
        timeout: float | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """
        Erstellt einen neuen ausstehenden Request.

        Args:
            timeout: Timeout in Sekunden
            metadata: Zusätzliche Metadaten

        Returns:
            Korrelations-ID

        Raises:
            RuntimeError: Wenn Maximum erreicht
        """
        async with self._lock:
            if len(self._pending) >= self.config.max_pending:
                raise RuntimeError(
                    f"Maximum ausstehende Requests erreicht: "
                    f"{self.config.max_pending}"
                )

            correlation_id = self._generate_id()

            # Stelle sicher dass ID eindeutig ist
            while correlation_id in self._pending:
                correlation_id = self._generate_id()

            request = PendingRequest[T](
                correlation_id=correlation_id,
                timeout=timeout or self.config.default_timeout,
                metadata=metadata or {},
                _future=asyncio.get_event_loop().create_future(),
            )

            self._pending[correlation_id] = request

        return correlation_id

    async def wait_for_response(
        self,
        correlation_id: str,
        timeout: float | None = None,
    ) -> T:
        """
        Wartet auf Antwort für einen Request.

        Args:
            correlation_id: Korrelations-ID
            timeout: Optional: Override Timeout

        Returns:
            Response-Daten

        Raises:
            KeyError: Request nicht gefunden
            TimeoutError: Timeout erreicht
            Exception: Bei Fehler im Request
        """
        request = self._pending.get(correlation_id)
        if not request:
            raise KeyError(f"Request nicht gefunden: {correlation_id}")

        effective_timeout = timeout or request.timeout

        try:
            result = await asyncio.wait_for(
                request._future,
                timeout=effective_timeout,
            )
            return result

        except asyncio.TimeoutError:
            async with self._lock:
                if correlation_id in self._pending:
                    request = self._pending[correlation_id]
                    request.state = RequestState.TIMEOUT
                    if request._future and not request._future.done():
                        request._future.cancel()

            raise TimeoutError(
                f"Timeout für Request {correlation_id} "
                f"nach {effective_timeout}s"
            )

    def complete(
        self,
        correlation_id: str,
        result: T,
    ) -> bool:
        """
        Markiert Request als abgeschlossen mit Ergebnis.

        Args:
            correlation_id: Korrelations-ID
            result: Ergebnis/Antwort

        Returns:
            True wenn Request gefunden und abgeschlossen
        """
        request = self._pending.get(correlation_id)
        if not request:
            self.logger.warning(
                f"Unbekannte Korrelations-ID: {correlation_id}"
            )
            return False

        if request.is_complete:
            self.logger.warning(
                f"Request bereits abgeschlossen: {correlation_id}"
            )
            return False

        request.state = RequestState.COMPLETED
        request.result = result

        if request._future and not request._future.done():
            request._future.set_result(result)

        return True

    def fail(
        self,
        correlation_id: str,
        error: Exception,
    ) -> bool:
        """
        Markiert Request als fehlgeschlagen.

        Args:
            correlation_id: Korrelations-ID
            error: Fehler

        Returns:
            True wenn Request gefunden
        """
        request = self._pending.get(correlation_id)
        if not request:
            return False

        if request.is_complete:
            return False

        request.state = RequestState.ERROR
        request.error = error

        if request._future and not request._future.done():
            request._future.set_exception(error)

        return True

    def cancel(self, correlation_id: str) -> bool:
        """
        Bricht einen Request ab.

        Args:
            correlation_id: Korrelations-ID

        Returns:
            True wenn abgebrochen
        """
        request = self._pending.get(correlation_id)
        if not request:
            return False

        if request.is_complete:
            return False

        request.state = RequestState.CANCELLED

        if request._future and not request._future.done():
            request._future.cancel()

        return True

    async def remove(self, correlation_id: str) -> PendingRequest[T] | None:
        """
        Entfernt einen Request.

        Args:
            correlation_id: Korrelations-ID

        Returns:
            Entfernter Request oder None
        """
        async with self._lock:
            return self._pending.pop(correlation_id, None)

    def get(self, correlation_id: str) -> PendingRequest[T] | None:
        """Holt Request nach ID."""
        return self._pending.get(correlation_id)

    def has_pending(self, correlation_id: str) -> bool:
        """Prüft ob Request existiert."""
        return correlation_id in self._pending

    async def cleanup_expired(self) -> int:
        """
        Entfernt abgelaufene Requests.

        Returns:
            Anzahl entfernter Requests
        """
        to_remove = []

        async with self._lock:
            for cid, request in self._pending.items():
                if request.is_expired():
                    to_remove.append(cid)
                    if request._future and not request._future.done():
                        request._future.cancel()

            for cid in to_remove:
                self._pending.pop(cid, None)

        if to_remove:
            self.logger.debug(f"Cleanup: {len(to_remove)} abgelaufene Requests entfernt")

        return len(to_remove)

    async def cancel_all(self) -> int:
        """
        Bricht alle ausstehenden Requests ab.

        Returns:
            Anzahl abgebrochener Requests
        """
        count = 0

        async with self._lock:
            for cid, request in self._pending.items():
                if not request.is_complete:
                    request.state = RequestState.CANCELLED
                    if request._future and not request._future.done():
                        request._future.cancel()
                    count += 1

        return count

    def get_pending_ids(self) -> list[str]:
        """Listet alle ausstehenden Request-IDs."""
        return [
            cid for cid, req in self._pending.items()
            if req.state == RequestState.PENDING
        ]

    def get_stats(self) -> dict[str, Any]:
        """Liefert Statistiken."""
        states = {}
        for request in self._pending.values():
            state_name = request.state.name
            states[state_name] = states.get(state_name, 0) + 1

        return {
            "total_pending": self.pending_count,
            "states": states,
            "oldest_request_age": max(
                (r.elapsed_seconds for r in self._pending.values()),
                default=0.0,
            ),
        }

    def __contains__(self, correlation_id: str) -> bool:
        return correlation_id in self._pending

    def __len__(self) -> int:
        return self.pending_count