"""InfraLens Python example: rollout capacity and policy freshness ledger."""

from dataclasses import dataclass


@dataclass(frozen=True)
class RolloutPlan:
    actors: int
    tokens_per_second: float
    interval_seconds: float
    learner_tokens_per_update: int


def capacity(plan: RolloutPlan) -> dict[str, float]:
    generated = plan.actors * plan.tokens_per_second * plan.interval_seconds
    backlog = generated - plan.learner_tokens_per_update
    return {
        "generated_tokens": generated,
        "learner_demand": float(plan.learner_tokens_per_update),
        "backlog_delta": backlog,
        "production_ratio": generated / plan.learner_tokens_per_update,
    }


def accept_rollout(learner_version: int, rollout_version: int, max_lag: int = 1) -> bool:
    return 0 <= learner_version - rollout_version <= max_lag


def smoke_test() -> None:
    result = capacity(RolloutPlan(actors=16, tokens_per_second=100, interval_seconds=60, learner_tokens_per_update=80_000))
    assert result["generated_tokens"] == 96_000
    assert result["backlog_delta"] == 16_000
    assert accept_rollout(learner_version=9, rollout_version=8)
    assert not accept_rollout(learner_version=9, rollout_version=6)
    print("rl_rollout_capacity_toy.py ok", result)


if __name__ == "__main__":
    smoke_test()