src/estimate.py

"""Cost / size estimation from the no-fetch estimate job's output."""
from __future__ import annotations

import math
import re
from dataclasses import dataclass
from typing import Iterable, Optional

from . import config

_ESTIMATE_RE = re.compile(r"ESTIMATE\s+n_records=(\d+)\s+total_bytes=(\d+)")


@dataclass(frozen=True)
class Estimate:
    n_records: int
    total_bytes: int
    flavor: str

    @property
    def total_gib(self) -> float:
        return self.total_bytes / (1024 ** 3)

    @property
    def rec_per_s(self) -> int:
        return config.FLAVOR_REC_PER_S.get(self.flavor, config.FLAVOR_REC_PER_S[config.DEFAULT_FLAVOR])

    @property
    def fetch_seconds(self) -> float:
        if self.n_records <= 0:
            return 0.0
        return self.n_records / self.rec_per_s

    @property
    def n_processes(self) -> int:
        return config.compute_processes(self.n_records, self.flavor)

    @property
    def merge_seconds(self) -> float:
        # Only multi-process runs produce shards that need merging.
        if self.n_processes <= 1:
            return 0.0
        return config.MERGE_BASE_S + self.total_bytes / config.MERGE_THROUGHPUT_BYTES_PER_S

    @property
    def duration_seconds(self) -> float:
        # Billed time = container/install overhead + fetch + (merge if multi-process).
        return config.CONTAINER_OVERHEAD_S + self.fetch_seconds + self.merge_seconds

    @property
    def hourly_usd(self) -> float:
        return config.FLAVOR_HOURLY_USD.get(self.flavor, config.FLAVOR_HOURLY_USD[config.DEFAULT_FLAVOR])

    @property
    def cost_usd(self) -> float:
        return self.duration_seconds / 3600.0 * self.hourly_usd

    @property
    def n_output_files(self) -> int:
        if self.total_bytes <= 0:
            return 0
        return max(1, math.ceil(self.total_bytes / config.WARC_TARGET_SIZE))


def parse_estimate(logs: Iterable[str], flavor: str) -> Optional[Estimate]:
    """Find the last `ESTIMATE n_records=.. total_bytes=..` line in the job logs."""
    found: Optional[Estimate] = None
    for line in logs:
        m = _ESTIMATE_RE.search(line)
        if m:
            found = Estimate(int(m.group(1)), int(m.group(2)), flavor)
    return found


def _fmt_cost(usd: float) -> str:
    if usd <= 0:
        return "$0.00"
    if usd < 0.01:
        return "<$0.01"
    return f"${usd:.2f}"


def _fmt_duration(seconds: float) -> str:
    seconds = int(round(seconds))
    if seconds < 60:
        return f"{seconds}s"
    if seconds < 3600:
        return f"{seconds // 60}m {seconds % 60}s"
    return f"{seconds // 3600}h {(seconds % 3600) // 60}m"


def actual_cost_usd(running_secs: int, flavor: str) -> float:
    hourly = config.FLAVOR_HOURLY_USD.get(flavor, config.FLAVOR_HOURLY_USD[config.DEFAULT_FLAVOR])
    return max(0, running_secs) / 3600.0 * hourly


def format_cost_report(running_secs: int, flavor: str) -> str:
    """One-line actual-cost report shown after a job finishes (HF bills running time)."""
    return (
        f"💰 **Actual cost:** ~{_fmt_cost(actual_cost_usd(running_secs, flavor))} "
        f"({running_secs}s of runtime on `{flavor}`)"
    )


def format_estimate(est: Estimate) -> str:
    """Markdown summary shown to the user before launching the fetch job."""
    if est.n_records == 0:
        return (
            "### ⚠️ No matching records\n"
            "The index query matched **0 records** for these hostnames/domains in the "
            "selected crawl(s), so there is nothing to repackage.\n\n"
            "Things to check:\n"
            "- **Hostnames match exactly** (e.g. `www.example.com`). To include all "
            "subdomains of a site, use the **Registered domains** field instead "
            "(e.g. `example.com`).\n"
            "- Try a different / additional crawl."
        )
    runtime_parts = [f"~{_fmt_duration(config.CONTAINER_OVERHEAD_S)} setup",
                     f"~{_fmt_duration(est.fetch_seconds)} fetch"]
    if est.merge_seconds > 0:
        runtime_parts.append(f"~{_fmt_duration(est.merge_seconds)} merge ({est.n_processes} shards)")
    return (
        "### 💸 Cost estimate\n"
        f"- **Records to fetch:** {est.n_records:,}\n"
        f"- **Total data size:** {est.total_gib:.2f} GiB\n"
        f"- **Output:** ~{est.n_output_files} WARC file(s) (≤1 GiB each)\n"
        f"- **Hardware:** `{est.flavor}` (${est.hourly_usd:.2f}/hr, ~{est.rec_per_s:,} rec/s)\n"
        f"- **Estimated runtime:** ~{_fmt_duration(est.duration_seconds)} "
        f"({' + '.join(runtime_parts)})\n"
        f"- **Estimated cost:** **~{_fmt_cost(est.cost_usd)}**\n\n"
        "_Estimate assumes conservative benchmark throughput; actual cost is billed "
        "per second of job runtime. The estimate job itself costs a few cents._"
    )