from __future__ import annotations

from dataclasses import dataclass
import random
import time
from typing import Callable, Sequence, TypeVar


T = TypeVar("T")


@dataclass(frozen=True)
class ChaosPolicy:
    drop_rate: float = 0.0
    duplicate_rate: float = 0.0
    max_delay_s: float = 0.0
    seed: int = 0

    def __post_init__(self) -> None:
        if not (0.0 <= self.drop_rate <= 1.0):
            raise ValueError("drop_rate must be within [0.0, 1.0]")
        if not (0.0 <= self.duplicate_rate <= 1.0):
            raise ValueError("duplicate_rate must be within [0.0, 1.0]")
        if self.max_delay_s < 0.0:
            raise ValueError("max_delay_s must be >= 0.0")


def dispatch_events_with_policy(
    events: Sequence[T],
    send: Callable[[T], None],
    *,
    policy: ChaosPolicy,
    allow_duplicate: Callable[[T], bool] | None = None,
) -> int:
    rng = random.Random(policy.seed)
    sent_count = 0
    for event in events:
        if policy.drop_rate > 0.0 and rng.random() < policy.drop_rate:
            continue
        send(event)
        sent_count += 1
        if policy.max_delay_s > 0.0:
            time.sleep(rng.uniform(0.0, policy.max_delay_s))

        if (
            policy.duplicate_rate > 0.0
            and (allow_duplicate(event) if allow_duplicate is not None else True)
            and rng.random() < policy.duplicate_rate
        ):
            send(event)
            sent_count += 1
            if policy.max_delay_s > 0.0:
                time.sleep(rng.uniform(0.0, policy.max_delay_s))
    return sent_count


__all__ = ["ChaosPolicy", "dispatch_events_with_policy"]
