from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Generic, TypeVar


T = TypeVar("T")


@dataclass(frozen=True)
class ThresholdProbe(Generic[T]):
    threshold: float
    feasible: bool
    payload: T


@dataclass(frozen=True)
class ThresholdSearchResult(Generic[T]):
    best_threshold: float
    best_feasible_payload: T | None
    probes: list[ThresholdProbe[T]]


def binary_search_max_feasible(
    *,
    low: float,
    high: float,
    tolerance: float,
    max_probes: int,
    evaluator: Callable[[float], ThresholdProbe[T]],
) -> ThresholdSearchResult[T]:
    probes: list[ThresholdProbe[T]] = []
    cache: dict[float, ThresholdProbe[T]] = {}
    best_threshold = low
    best_payload: T | None = None
    cur_low = low
    cur_high = high
    for _ in range(max_probes):
        if cur_high - cur_low <= tolerance and best_payload is not None:
            break
        # Thresholds are also trace-selection boundaries.  Rounding a midpoint
        # can cross a real sampling_u value and silently replay a different
        # workload than the materialized window specifies.
        threshold = (cur_low + cur_high) / 2.0
        probe = cache.get(threshold)
        if probe is None:
            probe = evaluator(threshold)
            cache[threshold] = probe
            probes.append(probe)
        if probe.feasible:
            if threshold >= best_threshold:
                best_threshold = threshold
                best_payload = probe.payload
            cur_low = threshold
        else:
            cur_high = threshold
    return ThresholdSearchResult(
        best_threshold=best_threshold,
        best_feasible_payload=best_payload,
        probes=probes,
    )
