"""slo_burn.py: SLO error budget and multi-window multi-burn-rate alert calculator.

Input: an SLO (target, window in days) and a per-minute series of (good, total) counts.
Output: SLI, error budget in events and minutes, budget consumed, burn rate per sliding
window, the page and ticket alert state per minute (SRE workbook chapter 5 thresholds),
time to exhaustion, and the Prometheus rule YAML built from the same numbers.
Standard library plus numpy and yaml. Run: python3 slo_burn.py
"""

from dataclasses import dataclass

import numpy as np
import yaml

MINUTES = {"5m": 5, "30m": 30, "1h": 60, "2h": 120, "6h": 360, "1d": 1440, "3d": 4320}

# SRE workbook, chapter 5, "6: Multiwindow, Multi-Burn-Rate Alerts". Each rule fires when the
# burn rate is above the threshold over BOTH the long and the short window (short = long / 12).
# (long, short, burn rate, budget consumed over the long window, severity)
RULES = [
    ("1h", "5m", 14.4, "2%", "page"),
    ("6h", "30m", 6.0, "5%", "page"),
    ("1d", "2h", 3.0, "10%", "ticket"),
    ("3d", "6h", 1.0, "10%", "ticket"),
]


@dataclass(frozen=True)
class SLO:
    target: float  # 0.999 means 99.9 percent
    window_days: int  # 30

    @property
    def window_minutes(self) -> int:
        return self.window_days * 1440

    @property
    def error_fraction(self) -> float:
        return 1.0 - self.target  # 0.001 for 99.9 percent


class Series:
    """Per-minute (good, total) counts. Index 0 is the oldest minute."""

    def __init__(self, good, total):
        self.good = np.asarray(good, dtype=np.int64)
        self.total = np.asarray(total, dtype=np.int64)
        if self.good.shape != self.total.shape or (self.good > self.total).any():
            raise ValueError("good and total must align and good <= total")
        self.bad = self.total - self.good
        # Prefix sums with a leading zero so any window sum is two lookups: O(1) per window.
        self.cum_bad = np.concatenate([[0], np.cumsum(self.bad)])
        self.cum_total = np.concatenate([[0], np.cumsum(self.total)])

    def __len__(self):
        return len(self.total)


# ---------- budget arithmetic ----------


def sli(s: Series) -> float:
    """Fraction of good events over the whole series (the SLI)."""
    return float(s.good.sum() / s.total.sum()) if s.total.sum() else 1.0


def error_budget_events(slo: SLO, s: Series) -> float:
    """How many bad events the SLO allows over the series: total x (1 - target)."""
    return float(s.total.sum()) * slo.error_fraction


def error_budget_minutes(slo: SLO) -> float:
    """Minutes of total outage the SLO allows per window: 43.2 for 99.9 percent over 30 days."""
    return slo.window_minutes * slo.error_fraction


def budget_consumed(slo: SLO, s: Series) -> float:
    """Fraction of the error budget spent so far (1.0 means exhausted, above 1 means breached)."""
    budget = error_budget_events(slo, s)
    return float(s.bad.sum()) / budget if budget else 0.0


# ---------- burn rate ----------


def burn_rate_series(slo: SLO, s: Series, window: str) -> np.ndarray:
    """Burn rate at every minute over a trailing window: error ratio in the window / (1 - target).

    Burn rate 1 means the budget is spent exactly at the end of the SLO window.
    Minutes before a full window exists use the partial window from the start, like rate().
    """
    w = MINUTES[window]
    n = len(s)
    end = np.arange(1, n + 1)
    start = np.maximum(end - w, 0)
    bad = s.cum_bad[end] - s.cum_bad[start]
    total = s.cum_total[end] - s.cum_total[start]
    with np.errstate(divide="ignore", invalid="ignore"):
        ratio = np.where(total > 0, bad / np.maximum(total, 1), 0.0)
    return ratio / slo.error_fraction


def burn_rates_now(slo: SLO, s: Series) -> dict:
    """Burn rate over every window at the last minute of the series."""
    return {w: float(burn_rate_series(slo, s, w)[-1]) for w in MINUTES}


def evaluate_alerts(slo: SLO, s: Series) -> dict:
    """Per-minute boolean arrays: page and ticket, each the OR of its two window pairs."""
    rates = {w: burn_rate_series(slo, s, w) for w in MINUTES}
    state = {"page": np.zeros(len(s), dtype=bool), "ticket": np.zeros(len(s), dtype=bool)}
    for long_w, short_w, rate, _, severity in RULES:
        state[severity] |= (rates[long_w] > rate) & (rates[short_w] > rate)
    return state


def time_to_exhaustion_minutes(slo: SLO, s: Series, window: str = "1h") -> float:
    """Minutes until the REMAINING budget is gone if the current burn rate (over `window`) holds.

    At burn rate b a full budget lasts window_minutes / b; only the unspent part is left.
    Returns inf when nothing is burning.
    """
    b = burn_rate_series(slo, s, window)[-1]
    remaining = max(0.0, 1.0 - budget_consumed(slo, s))
    return float("inf") if b <= 0 else slo.window_minutes * remaining / b


# ---------- Prometheus rules built from the same numbers ----------


def prometheus_rules(slo: SLO, job: str = "myjob", metric: str = "slo") -> str:
    """Recording rules per window plus the two multiwindow alerts, as a YAML string."""
    e = slo.error_fraction
    lines = ["groups:", "  - name: slo-burn-rate", "    rules:"]
    for w in MINUTES:
        lines += [
            f"      - record: job:{metric}_errors_per_request:ratio_rate{w}",
            "        expr: |",
            f"          sum(rate({metric}_errors[{w}])) by (job)",
            "          /",
            f"          sum(rate({metric}_requests[{w}])) by (job)",
        ]
    for severity in ("page", "ticket"):
        pairs = [r for r in RULES if r[4] == severity]
        clauses = []
        for long_w, short_w, rate, _, _ in pairs:
            clauses.append(
                f"          (\n"
                f"            job:{metric}_errors_per_request:ratio_rate{long_w}{{job=\"{job}\"}} > ({rate:g}*{e:g})\n"
                f"          and\n"
                f"            job:{metric}_errors_per_request:ratio_rate{short_w}{{job=\"{job}\"}} > ({rate:g}*{e:g})\n"
                f"          )"
            )
        lines += [
            f"      - alert: {'FastBurn' if severity == 'page' else 'SlowBurn'}ErrorBudget",
            "        expr: |",
            "\n          or\n".join(clauses),
            "        labels:",
            f"          severity: {severity}",
            "        annotations:",
            f"          summary: '{job} burns its {slo.target:.1%} / {slo.window_days}d error budget too fast'",
        ]
    return "\n".join(lines) + "\n"


# ---------- report ----------


def report(slo: SLO, s: Series) -> None:
    state = evaluate_alerts(slo, s)
    print(f"SLO {slo.target:.1%} over {slo.window_days} days ({slo.window_minutes} minutes)")
    print(f"  SLI so far            : {sli(s):.5%}")
    print(f"  error budget (events) : {error_budget_events(slo, s):,.0f}")
    print(f"  error budget (minutes): {error_budget_minutes(slo):.1f} minutes of full outage")
    print(f"  bad events so far     : {int(s.bad.sum()):,}")
    print(f"  budget consumed       : {budget_consumed(slo, s):.2%}")
    print("  burn rate now         : " + ", ".join(f"{w}={b:.2f}" for w, b in burn_rates_now(slo, s).items()))
    tte = time_to_exhaustion_minutes(slo, s)
    print(f"  time to exhaustion    : {'never at this rate' if tte == float('inf') else f'{tte / 60:.1f} hours'}")
    for sev in ("page", "ticket"):
        mins = int(state[sev].sum())
        print(f"  {sev:6s} alert          : {'FIRING' if state[sev][-1] else 'ok'} now, {mins} firing minutes in the series")


# ---------- self-checks ----------


def generate_series(slo: SLO, seed: int, per_minute: int, baseline_error: float,
                    incident_start: int, incident_minutes: int, incident_error: float) -> Series:
    """A full SLO window of traffic at `per_minute` requests, a small random baseline error
    rate, and one incident with an exact error rate (so the arithmetic is hand-checkable)."""
    rng = np.random.default_rng(seed)
    n = slo.window_minutes
    total = np.full(n, per_minute, dtype=np.int64)
    bad = rng.binomial(per_minute, baseline_error, size=n)
    bad[incident_start:incident_start + incident_minutes] = int(per_minute * incident_error)
    return Series(total - bad, total)


def main() -> None:
    slo = SLO(target=0.999, window_days=30)

    # Worked example from the workbook: burn 14.4 over 1 hour spends 14.4 * 60 / 43200 = 2 percent.
    assert abs(14.4 * 60 / slo.window_minutes - 0.02) < 1e-12
    assert abs(6 * 360 / slo.window_minutes - 0.05) < 1e-12
    assert abs(1 * 4320 / slo.window_minutes - 0.10) < 1e-12
    assert abs(3 * 1440 / slo.window_minutes - 0.10) < 1e-12
    assert abs(error_budget_minutes(slo) - 43.2) < 1e-9
    print("check ok: 14.4x for 1h = 2%, 6x for 6h = 5%, 3x for 1d = 10%, 1x for 3d = 10%; budget = 43.2 minutes")

    per_minute, baseline, inc_start, inc_len, inc_err = 1000, 0.0005, 20 * 1440 + 600, 20, 0.05
    s = generate_series(slo, seed=7, per_minute=per_minute, baseline_error=baseline,
                        incident_start=inc_start, incident_minutes=inc_len, incident_error=inc_err)
    report(slo, s)

    # Budget arithmetic against a hand-computed number: 43,200 minutes x 1,000 = 43.2 M events,
    # budget = 43,200 events; the incident alone is 20 x 50 = 1,000 bad events = 2.3148% of budget.
    assert abs(error_budget_events(slo, s) - 43_200) < 1e-6
    incident_bad = int(s.bad[inc_start:inc_start + inc_len].sum())
    assert incident_bad == 1000
    baseline_bad = int(s.bad.sum()) - incident_bad
    expected = (baseline_bad + 1000) / 43200
    assert abs(budget_consumed(slo, s) - expected) < 1e-12
    assert abs(1000 / 43200 - 0.023148) < 1e-6
    print(f"check ok: budget consumed {budget_consumed(slo, s):.4%} = ({baseline_bad} baseline + 1000 incident) / 43200")

    # Burn rate over a 5m window fully inside the incident is exactly 0.05 / 0.001 = 50.
    b5 = burn_rate_series(slo, s, "5m")
    assert abs(b5[inc_start + inc_len - 1] - 50.0) < 1e-9
    print("check ok: 5m burn rate at the end of the incident is exactly 50.0 (5% errors / 0.1% allowed)")

    # The fast-burn page fires during the incident and clears after it; the slow-burn ticket never fires.
    state = evaluate_alerts(slo, s)
    page, ticket = state["page"], state["ticket"]
    first = int(np.argmax(page))
    assert page.any() and inc_start <= first < inc_start + inc_len
    assert page[inc_start + inc_len - 1]
    cleared = inc_start + inc_len + int(np.argmin(page[inc_start + inc_len:]))
    assert not page[cleared] and cleared - (inc_start + inc_len) <= 5
    assert not page[cleared:].any() and not page[:inc_start].any()
    assert not ticket.any()
    print(f"check ok: page fired {first - inc_start} minutes into the incident, "
          f"cleared {cleared - (inc_start + inc_len)} minutes after it ended (short window reset)")
    full_3d = burn_rate_series(slo, s, "3d")[MINUTES["3d"]:]  # ignore partial windows at the start
    print(f"check ok: ticket (slow burn) never fired: 3d burn peaked at {full_3d.max():.2f}, below 1.0")

    # Time to exhaustion: 30 clean days then one hour at 1.44% errors is burn 14.4 over 1h.
    # A full budget lasts 720 h / 14.4 = 50 h; the hour already spent about 2%, so a bit less is left.
    n = slo.window_minutes
    hot = Series(good=[10000] * n + [10000 - 144] * 60, total=[10000] * (n + 60))
    assert abs(burn_rate_series(slo, hot, "1h")[-1] - 14.4) < 1e-9
    tte = time_to_exhaustion_minutes(slo, hot, "1h")
    left = 1 - budget_consumed(slo, hot)
    assert abs(tte / 60 - 720 * left / 14.4) < 1e-9 and 49.0 < tte / 60 < 50.0
    print(f"check ok: at 14.4x burn a full budget lasts 50.0 h; with {left:.2%} left it lasts {tte / 60:.1f} h")

    # The YAML parses, and the thresholds inside it are the rule numbers times (1 - target).
    text = prometheus_rules(slo)
    doc = yaml.safe_load(text)
    rules = doc["groups"][0]["rules"]
    records = [r for r in rules if "record" in r]
    alerts = [r for r in rules if "alert" in r]
    assert len(records) == len(MINUTES) and len(alerts) == 2
    assert "(14.4*0.001)" in alerts[0]["expr"] and "ratio_rate5m" in alerts[0]["expr"]
    assert "(1*0.001)" in alerts[1]["expr"] and "ratio_rate3d" in alerts[1]["expr"]
    assert alerts[0]["labels"]["severity"] == "page" and alerts[1]["labels"]["severity"] == "ticket"
    print(f"check ok: Prometheus YAML parses with {len(records)} recording rules and 2 alerts")
    print("\n--- prometheus rules ---")
    print(text)
    print("all checks passed")


if __name__ == "__main__":
    main()