"""Lab 01 — Platform Reference Model & Budget Calculator.

Fill in every ``# TODO``. Work top to bottom: section 1 is used by section 2, and the
admission pipeline in section 7 is the synthesis.

Rules (see ../../LAB-STANDARD.md):
  * pure stdlib, deterministic — no wall clock, no randomness, no I/O
  * money is integer micro-USD so accumulated cost is exact
  * validate inputs and raise ``ValueError`` / ``KeyError`` rather than returning junk

Run:
    pytest test_lab.py -v
    LAB_MODULE=solution pytest test_lab.py -v    # the reference, must be green
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

MINUTES_PER_DAY = 24 * 60

# Budgets and burn rates are ratios of floats, so an exact threshold comparison is
# unreliable: ``1 - 0.999`` is 0.0009999999999998899, not 0.001, and a burn rate that is
# mathematically exactly 14.4 computes as 14.399999999999986. Real alerting systems hit
# this and silently fail to fire at the boundary. Use this tolerance in
# ``BudgetLedger.policy_state`` and ``MultiWindowAlertPolicy.evaluate``.
EPSILON = 1e-9


def _check_probability(value: float, label: str) -> None:
    if not 0.0 <= value <= 1.0:
        raise ValueError(f"{label} must be in [0, 1], got {value!r}")


# --------------------------------------------------------------------------------------
# 1. Availability composition
# --------------------------------------------------------------------------------------


def series_availability(values: Sequence[float]) -> float:
    """Availability of components that must ALL succeed.

    Returns the product. An EMPTY chain returns 1.0 (the identity of a product).
    Validate each value with ``_check_probability``.
    """
    # TODO: multiply the values, validating each one; empty -> 1.0
    raise NotImplementedError


def parallel_availability(values: Sequence[float]) -> float:
    """Availability of redundant components: ``1 - prod(1 - A_i)``.

    An EMPTY group returns 0.0 — no members means nothing can serve.
    """
    # TODO: multiply the unavailabilities and subtract from 1; empty -> 0.0
    raise NotImplementedError


def correlated_parallel_availability(values: Sequence[float], common_mode: float) -> float:
    """Redundancy with a shared failure mode.

        unavailability = c * mean(u_i) + (1 - c) * prod(u_i)      where u_i = 1 - A_i

    ``c = 0`` reduces to :func:`parallel_availability`; ``c = 1`` means redundancy buys
    nothing. Empty group -> 0.0.
    """
    # TODO: implement the mixture above; validate common_mode and each availability
    raise NotImplementedError


# --------------------------------------------------------------------------------------
# 2. The platform model
# --------------------------------------------------------------------------------------


@dataclass(frozen=True)
class Component:
    """One dependency in the request path.

    ``degradable=True`` means the request still succeeds when this component fails, at
    reduced quality — so it does NOT multiply into request availability.
    """

    name: str
    availability: float
    degradable: bool = False
    p95_ms: int = 0

    def __post_init__(self) -> None:
        _check_probability(self.availability, f"availability of {self.name!r}")
        if self.p95_ms < 0:
            raise ValueError("p95_ms must be >= 0")


@dataclass(frozen=True)
class PlatformModel:
    name: str
    components: Tuple[Component, ...]

    @classmethod
    def of(cls, name: str, components: Iterable[Component]) -> "PlatformModel":
        return cls(name=name, components=tuple(components))

    def request_availability(self) -> float:
        """Product over NON-degradable components only."""
        # TODO
        raise NotImplementedError

    def quality_availability(self) -> float:
        """Product over ALL components."""
        # TODO
        raise NotImplementedError

    def downtime_minutes(self, window_days: int = 30, *, quality: bool = False) -> float:
        """``(1 - A) * window_days * MINUTES_PER_DAY``, using the quality availability
        when ``quality=True``."""
        # TODO
        raise NotImplementedError

    def weakest_links(self, k: int = 3, *, include_degradable: bool = False) -> List[Tuple[str, float]]:
        """Top-k components by unavailability, as ``(name, 1 - A)``.

        Sort by unavailability DESC, breaking ties on name ASC so the result is
        deterministic. Degradable components are excluded unless asked for.
        """
        # TODO
        raise NotImplementedError

    def with_degradable(self, *names: str) -> "PlatformModel":
        """Copy with the named components marked degradable. Unknown name -> KeyError."""
        # TODO
        raise NotImplementedError

    def with_replaced(self, name: str, availability: float) -> "PlatformModel":
        """Copy with one component's availability replaced. Unknown name -> KeyError."""
        # TODO
        raise NotImplementedError


# --------------------------------------------------------------------------------------
# 3. Error budgets
# --------------------------------------------------------------------------------------


@dataclass(frozen=True)
class ErrorBudget:
    slo: float
    window_days: int = 30

    def __post_init__(self) -> None:
        _check_probability(self.slo, "slo")
        if self.window_days <= 0:
            raise ValueError("window_days must be > 0")

    def total_minutes(self) -> float:
        """``(1 - slo) * window_days * MINUTES_PER_DAY``."""
        # TODO
        raise NotImplementedError

    def allocate(self, shares: Mapping[str, float], *, tolerance: float = 1e-9) -> Dict[str, float]:
        """Split the budget by weight.

        Raise ``ValueError`` for an empty mapping, a negative share, or weights that do
        not sum to 1.0 within ``tolerance``.
        """
        # TODO
        raise NotImplementedError


class BudgetLedger:
    """Tracks consumption of an allocated budget. NEVER reports a negative remainder —
    overspend is surfaced separately by :meth:`overspend_for`."""

    def __init__(self, budget: ErrorBudget, shares: Mapping[str, float]) -> None:
        self.budget = budget
        self.allocations = budget.allocate(shares)
        self._consumed: Dict[str, float] = {layer: 0.0 for layer in self.allocations}

    def consume(self, layer: str, minutes: float) -> None:
        """Unknown layer -> KeyError. Negative minutes -> ValueError."""
        # TODO
        raise NotImplementedError

    def consumed_for(self, layer: str) -> float:
        return self._consumed[layer]

    def remaining_for(self, layer: str) -> float:
        """``max(0, allocation - consumed)``."""
        # TODO
        raise NotImplementedError

    def overspend_for(self, layer: str) -> float:
        """``max(0, consumed - allocation)``."""
        # TODO
        raise NotImplementedError

    def total_consumed(self) -> float:
        return sum(self._consumed.values())

    def remaining(self) -> float:
        """``max(0, total budget - total consumed)``."""
        # TODO
        raise NotImplementedError

    def fraction_remaining(self) -> float:
        """``remaining / total``; 0.0 when the total is 0."""
        # TODO
        raise NotImplementedError

    def policy_state(self) -> str:
        """The written error-budget policy, as a pure function of fraction remaining:

            <= EPSILON -> "freeze"
            <  0.25    -> "reliability-focus"
            <  0.50    -> "elevated"
            else       -> "normal"

        Use EPSILON, not 0.0: a budget consumed exactly to the minute leaves a float
        residue of ~1e-14, and reporting "reliability-focus" when the budget is gone is
        the kind of bug that only shows up during an incident.
        """
        # TODO
        raise NotImplementedError


def burn_rate(observed_bad_ratio: float, slo: float) -> float:
    """``observed_bad_ratio / (1 - slo)``.

    A perfect SLO (1.0) has no allowance: return ``math.inf`` when any error is
    observed, else 0.0.
    """
    # TODO
    raise NotImplementedError


def burn_rate_threshold(budget_fraction: float, window_hours: float, period_days: int = 30) -> float:
    """Derive the threshold instead of memorizing 14.4.

        B = budget_fraction * (period_hours / window_hours)

    ``burn_rate_threshold(0.02, 1)`` -> 14.4;  ``burn_rate_threshold(0.05, 6)`` -> 6.0.
    """
    # TODO: validate window_hours > 0, period_days > 0, budget_fraction in [0, 1]
    raise NotImplementedError


@dataclass(frozen=True)
class AlertRule:
    name: str
    severity: str  # "page" | "ticket"
    long_window_hours: float
    short_window_hours: float
    threshold: float


DEFAULT_ALERT_RULES: Tuple[AlertRule, ...] = (
    AlertRule("fast-burn", "page", 1.0, 1.0 / 12, 14.4),
    AlertRule("medium-burn", "page", 6.0, 0.5, 6.0),
    AlertRule("slow-burn", "ticket", 72.0, 6.0, 1.0),
)


@dataclass(frozen=True)
class MultiWindowAlertPolicy:
    slo: float
    rules: Tuple[AlertRule, ...] = DEFAULT_ALERT_RULES

    def evaluate(self, bad_ratio_over: Callable[[float], float]) -> List[AlertRule]:
        """Fire a rule only when BOTH its windows are at or above the threshold.

        Preserve rule order in the returned list. Compare against
        ``threshold * (1 - EPSILON)`` so a burn rate that is mathematically exactly at
        the threshold fires — that boundary is a favourite interview probe, and a naive
        ``>=`` on raw floats misses it.
        """
        # TODO
        raise NotImplementedError

    def highest_severity(self, fired: Sequence[AlertRule]) -> Optional[str]:
        """"page" if any fired rule pages, else "ticket" if any fired, else None."""
        # TODO
        raise NotImplementedError


# --------------------------------------------------------------------------------------
# 4. Latency budgets
# --------------------------------------------------------------------------------------


@dataclass(frozen=True)
class LatencyStage:
    name: str
    p95_ms: int
    group: Optional[str] = None  # stages sharing a group run concurrently
    sheddable: bool = False

    def __post_init__(self) -> None:
        if self.p95_ms < 0:
            raise ValueError("p95_ms must be >= 0")


@dataclass(frozen=True)
class LatencyBudget:
    target_p95_ms: int
    stages: Tuple[LatencyStage, ...]

    @classmethod
    def of(cls, target_p95_ms: int, stages: Iterable[LatencyStage]) -> "LatencyBudget":
        if target_p95_ms <= 0:
            raise ValueError("target_p95_ms must be > 0")
        return cls(target_p95_ms=target_p95_ms, stages=tuple(stages))

    def committed_ms(self, *, shed: Sequence[str] = ()) -> int:
        """Sum of ungrouped stages plus, for each parallel group, its MAX.

        Stages named in ``shed`` are skipped entirely.
        """
        # TODO
        raise NotImplementedError

    def headroom_ms(self, *, shed: Sequence[str] = ()) -> int:
        # TODO
        raise NotImplementedError

    def is_feasible(self, *, shed: Sequence[str] = ()) -> bool:
        # TODO
        raise NotImplementedError

    def fits_fallback(self, fallback_timeout_ms: int, *, shed: Sequence[str] = ()) -> bool:
        """A fallback that does not fit the remaining budget is decoration."""
        # TODO: negative timeout -> ValueError
        raise NotImplementedError

    def shed_order(self) -> List[str]:
        """The degradation ladder: sheddable stage names, most expensive first, ties on
        name ascending."""
        # TODO
        raise NotImplementedError

    def shed_until_fits(self, fallback_timeout_ms: int) -> List[str]:
        """Shed in ladder order until the fallback fits; return the stages shed.

        May return the whole ladder without succeeding — the caller re-checks.
        """
        # TODO
        raise NotImplementedError


# --------------------------------------------------------------------------------------
# 5. Loop reliability
# --------------------------------------------------------------------------------------


def loop_success(p: float, n: int) -> float:
    """``p ** n``. Negative ``n`` -> ValueError."""
    # TODO
    raise NotImplementedError


def effective_step_probability(p: float, retries: int) -> float:
    """``1 - (1 - p) ** (retries + 1)``. Negative retries -> ValueError."""
    # TODO
    raise NotImplementedError


def max_steps_for_target(p: float, target: float) -> int:
    """Largest ``n`` with ``p ** n >= target``: ``floor(log(target) / log(p))``.

    ``p >= 1.0`` -> ValueError (no bound). ``target <= 0`` -> ValueError.
    ``p == 0`` -> 0.
    """
    # TODO
    raise NotImplementedError


# --------------------------------------------------------------------------------------
# 6. Cost
# --------------------------------------------------------------------------------------


@dataclass(frozen=True)
class TokenPrices:
    """Micro-USD per 1 000 tokens."""

    input_per_1k: int
    cached_input_per_1k: int
    output_per_1k: int

    def __post_init__(self) -> None:
        for label in ("input_per_1k", "cached_input_per_1k", "output_per_1k"):
            if getattr(self, label) < 0:
                raise ValueError(f"{label} must be >= 0")


@dataclass(frozen=True)
class CostModel:
    prices: TokenPrices

    def step_cost_micros(self, tokens_in: int, tokens_cached: int, tokens_out: int) -> int:
        """``((in - cached) * c_in + cached * c_cache + out * c_out) // 1000``.

        Divide LAST so integer truncation happens once. Negative counts -> ValueError;
        ``tokens_cached > tokens_in`` -> ValueError.
        """
        # TODO
        raise NotImplementedError

    def run_cost_micros(
        self,
        *,
        base_tokens: int,
        per_step_tokens: int,
        output_tokens_per_step: int,
        steps: int,
        cached_prefix_tokens: int = 0,
    ) -> int:
        """Cost of an n-step run with an accumulating scratchpad.

        Step ``i`` (1-indexed) sends ``base + per_step * (i - 1)`` input tokens, of which
        ``min(cached_prefix_tokens, tokens_in)`` are cached.
        """
        # TODO
        raise NotImplementedError

    @staticmethod
    def total_input_tokens(base_tokens: int, per_step_tokens: int, steps: int) -> int:
        """Closed form: ``steps * base + per_step * steps * (steps - 1) // 2``."""
        # TODO
        raise NotImplementedError

    @staticmethod
    def cost_per_successful_action_micros(attempt_cost_micros: int, success_probability: float) -> int:
        """``round(attempt_cost / p_success)``. ``p_success == 0`` -> ValueError."""
        # TODO
        raise NotImplementedError


def effective_cost_with_cache(miss_cost: float, hit_cost: float, hit_rate: float) -> float:
    """``(1 - h) * miss + h * hit``."""
    # TODO
    raise NotImplementedError


def cache_savings_fraction(miss_cost: float, hit_cost: float, hit_rate: float) -> float:
    """``h * (1 - hit / miss)``. ``miss_cost <= 0`` -> ValueError."""
    # TODO
    raise NotImplementedError


# --------------------------------------------------------------------------------------
# 7. The admission pipeline — which layer denies, and why
# --------------------------------------------------------------------------------------

LAYERS: Tuple[str, ...] = (
    "users_and_channels",
    "control_plane",
    "agent_kernel",
    "knowledge_foundation",
    "action_gateway",
)


@dataclass(frozen=True)
class AgentRegistration:
    agent_id: str
    tenant: str
    permitted_tools: Tuple[str, ...]
    granted_scopes: Tuple[str, ...]
    max_action_amount_micros: int
    evaluation_fresh: bool = True


@dataclass(frozen=True)
class ProposedAction:
    tenant: str  # derived from the verified token, NEVER from the request body
    agent_id: str
    tool: str
    channel: str
    user_authenticated: bool = True
    amount_micros: int = 0
    resource_tenant: str = ""
    approvals: Tuple[str, ...] = ()
    step_index: int = 1
    run_cost_micros: int = 0
    derived_from_untrusted_content: bool = False
    retrieved_tenants: Tuple[str, ...] = ()

    def __post_init__(self) -> None:
        if self.amount_micros < 0:
            raise ValueError("amount_micros must be >= 0")
        if self.step_index < 1:
            raise ValueError("step_index must be >= 1")


@dataclass(frozen=True)
class Denial:
    layer: str
    code: str
    reason: str


@dataclass(frozen=True)
class AdmissionResult:
    allowed: bool
    denials: Tuple[Denial, ...]

    @property
    def primary(self) -> Optional[Denial]:
        """The earliest denial — what the caller is told."""
        return self.denials[0] if self.denials else None

    @property
    def defence_depth(self) -> int:
        """Number of DISTINCT layers that denied. 1 means a single point of failure
        with good intentions."""
        return len({d.layer for d in self.denials})


@dataclass(frozen=True)
class AdmissionPipeline:
    """Runs EVERY layer's checks and reports ALL denials — not just the first.

    Denials are returned sorted by layer order (``LAYERS``) then by code, so the output
    is deterministic and diffable.
    """

    registry: Mapping[str, AgentRegistration]
    tool_required_scope: Mapping[str, str]
    side_effecting_tools: frozenset
    dual_control_threshold_micros: int
    approval_capable_channels: frozenset
    max_steps: int = 25
    max_run_cost_micros: int = 5_000_000

    def evaluate(self, action: ProposedAction) -> AdmissionResult:
        """Collect denials from all five layers, sort, and wrap in an AdmissionResult."""
        # TODO: call the five helpers, sort by (LAYERS index, code), build the result
        raise NotImplementedError

    def _channel_checks(self, action: ProposedAction) -> List[Denial]:
        """users_and_channels:
        - UNAUTHENTICATED         : no authenticated human principal
        - CHANNEL_CANNOT_APPROVE  : amount >= dual-control threshold but the channel is
                                    not in ``approval_capable_channels``
        """
        # TODO
        raise NotImplementedError

    def _control_plane_checks(
        self, action: ProposedAction, registration: Optional[AgentRegistration]
    ) -> List[Denial]:
        """control_plane:
        - AGENT_NOT_REGISTERED    : no registry entry (return immediately — the other
                                    checks have nothing to read)
        - TOOL_NOT_PERMITTED      : tool outside the registered tool set
        - EVALUATION_STALE        : KYA posture check failed
        - AGENT_TENANT_MISMATCH   : registry tenant != token tenant
        """
        # TODO
        raise NotImplementedError

    def _kernel_checks(self, action: ProposedAction) -> List[Denial]:
        """agent_kernel:
        - STEP_BUDGET_EXCEEDED    : step_index > max_steps
        - COST_CEILING_EXCEEDED   : run_cost_micros > max_run_cost_micros
        """
        # TODO
        raise NotImplementedError

    def _knowledge_checks(self, action: ProposedAction) -> List[Denial]:
        """knowledge_foundation:
        - CROSS_TENANT_RETRIEVAL     : any retrieved tenant != the caller's tenant
        - UNTRUSTED_INSTRUCTION_SOURCE : a side-effecting tool whose call was derived
                                       from retrieved content
        """
        # TODO
        raise NotImplementedError

    def _gateway_checks(
        self, action: ProposedAction, registration: Optional[AgentRegistration]
    ) -> List[Denial]:
        """action_gateway:
        - TENANT_MISMATCH         : resource_tenant set and != caller tenant
        - SCOPE_MISSING           : tool requires a scope the credential lacks
                                    (missing registration counts as missing scope)
        - ACTION_LIMIT_EXCEEDED   : amount above the agent's registered limit
        - DUAL_CONTROL_REQUIRED   : amount >= threshold and fewer than 2 DISTINCT
                                    approvers other than the agent itself
        """
        # TODO
        raise NotImplementedError


def main() -> None:
    """Optional: print your own worked example. See solution.py for the reference."""
    print("implement the TODOs, then compare with `python solution.py`")


if __name__ == "__main__":
    main()