"""System Design Lesson 1: small capacity and storage estimator.

This uses only the Python standard library.
"""

from __future__ import annotations

from dataclasses import dataclass

SECONDS_PER_DAY = 86_400


@dataclass(frozen=True)
class Workload:
    daily_active_users: int
    actions_per_user_per_day: float
    peak_factor: float
    bytes_per_write: int
    write_fraction: float
    retention_days: int
    replication_factor: int = 3

    def estimate(self) -> dict[str, float]:
        actions_per_day = self.daily_active_users * self.actions_per_user_per_day
        average_qps = actions_per_day / SECONDS_PER_DAY
        writes_per_day = actions_per_day * self.write_fraction
        logical_bytes = writes_per_day * self.bytes_per_write * self.retention_days
        return {
            "actions_per_day": actions_per_day,
            "average_qps": average_qps,
            "peak_qps": average_qps * self.peak_factor,
            "logical_storage_gb": logical_bytes / 1_000_000_000,
            "replicated_storage_gb": logical_bytes
            * self.replication_factor
            / 1_000_000_000,
        }


def main() -> None:
    workload = Workload(
        daily_active_users=10_000_000,
        actions_per_user_per_day=10,
        peak_factor=5,
        bytes_per_write=1_000,
        write_fraction=0.1,
        retention_days=365,
    )
    estimates = workload.estimate()
    assert round(estimates["average_qps"]) == 1_157
    for name, value in estimates.items():
        print(f"{name:24} {value:,.2f}")

    # Try it yourself:
    # 1. Model a write-heavy metrics pipeline.
    # 2. Add average response bytes and outbound bandwidth.
    # 3. Add one hot key receiving 20% of peak traffic.


if __name__ == "__main__":
    main()