app/estimates.py

from __future__ import annotations

import math

DEFAULT_SEGMENT_SECONDS = 60
DEFAULT_MAX_PARALLEL = 4
SECONDS_PER_SEGMENT_PRECISION = 180
SECONDS_PER_SEGMENT_SPEED = 90
OVERHEAD_SECONDS = 30
COST_PER_SEGMENT_USD = 0.12


def format_duration(seconds: int) -> str:
    if seconds < 60:
        return f"{seconds}s"
    minutes, secs = divmod(seconds, 60)
    if minutes < 60:
        return f"{minutes}m {secs}s" if secs else f"{minutes}m"
    hours, minutes = divmod(minutes, 60)
    return f"{hours}h {minutes}m"


def estimate_run(duration_seconds: float, mode: str = "precision") -> dict[str, object]:
    duration = max(1.0, float(duration_seconds))
    segments = max(1, math.ceil(duration / DEFAULT_SEGMENT_SECONDS))
    per_segment = SECONDS_PER_SEGMENT_PRECISION if mode == "precision" else SECONDS_PER_SEGMENT_SPEED
    batches = math.ceil(segments / DEFAULT_MAX_PARALLEL)
    wall_seconds = batches * per_segment + OVERHEAD_SECONDS
    cost_usd = round(segments * COST_PER_SEGMENT_USD, 2)
    return {
        "duration_seconds": round(duration, 1),
        "segments": segments,
        "segment_seconds": DEFAULT_SEGMENT_SECONDS,
        "max_parallel": DEFAULT_MAX_PARALLEL,
        "mode": mode,
        "estimated_time_seconds": wall_seconds,
        "estimated_time_label": format_duration(wall_seconds),
        "estimated_cost_usd": cost_usd,
        "disclaimer": "Rough estimate. Actual time and Replicate cost may vary.",
    }