"""Lesson 3 lab: range-based workload and capacity estimation.

Run:
    python 03-back-of-envelope-estimation.py

This lab uses only the Python standard library. Its numbers are assumptions,
not universal constants; change them and inspect which decision moves.
"""

from __future__ import annotations

from dataclasses import dataclass
from math import ceil

SECONDS_PER_DAY = 86_400
DECIMAL_GB = 1_000_000_000


@dataclass(frozen=True)
class Scenario:
    daily_active_users: int
    actions_per_user_day: float
    recipients_per_action: float
    peak_factor: float
    bytes_per_record: int
    index_multiplier: float
    retention_days: int
    replicas: int
    response_bytes: int
    cacheable_read_fraction: float
    cache_hit_rate: float
    safe_ops_per_worker: float
    target_utilization: float
    unavailable_workers: int

    def estimate(self) -> dict[str, float]:
        if not 0 < self.target_utilization <= 1:
            raise ValueError("target_utilization must be in (0, 1]")
        if not 0 <= self.cache_hit_rate <= 1:
            raise ValueError("cache_hit_rate must be in [0, 1]")

        product_actions = self.daily_active_users * self.actions_per_user_day
        operations = product_actions * self.recipients_per_action
        average_ops = operations / SECONDS_PER_DAY
        peak_ops = average_ops * self.peak_factor
        origin_peak_ops = peak_ops * (
            1 - self.cacheable_read_fraction * self.cache_hit_rate
        )
        logical_storage = operations * self.bytes_per_record * self.retention_days
        indexed_storage = logical_storage * self.index_multiplier
        physical_storage = indexed_storage * self.replicas
        peak_egress_bps = peak_ops * self.response_bytes * 8
        usable_worker_rate = self.safe_ops_per_worker * self.target_utilization
        workers_for_load = ceil(peak_ops / usable_worker_rate)
        provisioned_workers = workers_for_load + self.unavailable_workers

        return {
            "product_actions_day": product_actions,
            "operations_day": operations,
            "average_ops_s": average_ops,
            "peak_ops_s": peak_ops,
            "origin_peak_ops_s": origin_peak_ops,
            "logical_storage_gb": logical_storage / DECIMAL_GB,
            "indexed_storage_gb": indexed_storage / DECIMAL_GB,
            "physical_storage_gb": physical_storage / DECIMAL_GB,
            "peak_egress_mbps": peak_egress_bps / 1_000_000,
            "provisioned_workers": provisioned_workers,
        }


def estimate_range(low: Scenario, expected: Scenario, high: Scenario) -> None:
    estimates = [low.estimate(), expected.estimate(), high.estimate()]
    print(f"{'metric':24} {'low':>14} {'expected':>14} {'high':>14}")
    for metric in estimates[0]:
        values = [result[metric] for result in estimates]
        print(f"{metric:24} {values[0]:14,.2f} {values[1]:14,.2f} {values[2]:14,.2f}")


def main() -> None:
    base = dict(
        daily_active_users=20_000_000,
        actions_per_user_day=4,
        recipients_per_action=3,
        peak_factor=6,
        bytes_per_record=700,
        index_multiplier=1.4,
        retention_days=365,
        replicas=3,
        response_bytes=1_200,
        cacheable_read_fraction=0.8,
        cache_hit_rate=0.9,
        safe_ops_per_worker=800,
        target_utilization=0.65,
        unavailable_workers=2,
    )
    low = Scenario(**{**base, "daily_active_users": 12_000_000, "peak_factor": 4})
    expected = Scenario(**base)
    high = Scenario(
        **{**base, "daily_active_users": 32_000_000, "recipients_per_action": 5, "peak_factor": 10}
    )

    result = expected.estimate()
    assert round(result["average_ops_s"]) == 2_778
    assert result["peak_ops_s"] > result["origin_peak_ops_s"]
    assert result["indexed_storage_gb"] == 1.4 * result["logical_storage_gb"]
    assert result["physical_storage_gb"] == 3 * result["indexed_storage_gb"]
    assert result["provisioned_workers"] >= 35
    estimate_range(low, expected, high)

    # Try it yourself:
    # 1. Set cache_hit_rate to zero and quantify origin-load amplification.
    # 2. Model a celebrity key receiving 15% of peak traffic on one partition.
    # 3. Add compression, backups, and one year of 30% growth.


if __name__ == "__main__":
    main()