"""Runtime-cost tracking helpers for experiment stages."""

from __future__ import annotations

from datetime import datetime
from pathlib import Path

from abreka.persistence import load_experiment, save_experiment, update_index_entry

STAGE_ORDER: tuple[str, ...] = ("propose", "implement", "test", "run", "run_repair", "evaluate")


def _as_int(value: object, default: int = 0) -> int:
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


def _as_float(value: object, default: float = 0.0) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def stage_ordered_keys(runtime_costs: dict[str, dict[str, object]]) -> list[str]:
    """Return stage keys in canonical order, then unknown stages sorted."""
    present = set(runtime_costs)
    ordered = [stage for stage in STAGE_ORDER if stage in present]
    ordered.extend(sorted(present - set(STAGE_ORDER)))
    return ordered


def record_stage_runtime(
    abreka_dir: Path,
    exp_id: str,
    *,
    stage: str,
    elapsed_seconds: float,
    status: str,
    started_at: datetime,
    finished_at: datetime,
) -> None:
    """Record one stage attempt's runtime for an experiment.

    This accumulates attempts and total elapsed seconds per stage, while also
    storing last-attempt metadata for quick CLI display.
    """
    exp = load_experiment(abreka_dir, exp_id)
    runtime_costs = dict(exp.runtime_costs or {})

    previous_raw = runtime_costs.get(stage, {})
    previous: dict[str, object] = previous_raw if isinstance(previous_raw, dict) else {}

    elapsed = max(0.0, float(elapsed_seconds))
    attempts = _as_int(previous.get("attempts")) + 1
    total_seconds = _as_float(previous.get("seconds_total")) + elapsed

    stage_entry = dict(previous)
    stage_entry.update(
        {
            "attempts": attempts,
            "seconds_total": total_seconds,
            "seconds_last": elapsed,
            "last_status": status,
            "last_started_at": started_at.isoformat(),
            "last_finished_at": finished_at.isoformat(),
        }
    )

    runtime_costs[stage] = stage_entry
    exp.runtime_costs = runtime_costs

    save_experiment(exp, abreka_dir)
    update_index_entry(abreka_dir, exp)
