src/progress.py

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

# Exact worker-facing steps emitted by src/worker_payload.py.
# Keep this list close to the Job logs/events rather than a marketing pipeline.
STEP_ORDER = [
    "bootstrap",
    "dependencies",
    "auth",
    "model_analysis",
    "workspace",
    "node",
    "pi_install",
    "pi_config",
    "pi_run",
    "pi_verification",
    "metadata_sanitize",
    "requirements_sanitize",
    "hardware_strategy",
    "create_space_hardware",
    "create_space",
    "repair",
    "upload_files",
    "space_runtime",
    "space_logs",
    "api_validation",
    "live_wait",
    "generation_smoke",
    "inference_gate",
    "report_write",
    "done",
    "failure",
]

STEP_LABELS = {
    "bootstrap": "Bootstrap",
    "dependencies": "Dependencies",
    "auth": "Auth",
    "model_analysis": "Model analysis",
    "workspace": "Workspace",
    "node": "Node/npm",
    "pi_install": "Pi install",
    "pi_config": "Pi config",
    "pi_run": "Pi run",
    "pi_verification": "Pi verification",
    "metadata_sanitize": "Metadata sanitize",
    "requirements_sanitize": "Requirements sanitize",
    "hardware_strategy": "Hardware strategy",
    "create_space_hardware": "Create with hardware",
    "create_space": "Create Space",
    "repair": "Repair pass",
    "upload_files": "Upload files",
    "space_runtime": "Space runtime",
    "space_logs": "Space logs",
    "api_validation": "API validation",
    "live_wait": "Live wait",
    "generation_smoke": "Generation smoke",
    "inference_gate": "Inference gate",
    "report_write": "Report",
    "done": "Done",
    "failure": "Failure",
}

STEP_ALIASES = {
    "bucket_ready": "bootstrap",
    "job_launched": "bootstrap",
    "hardware_preferred": "hardware_strategy",
    "hardware_fallback": "hardware_strategy",
    "hardware": "hardware_strategy",
    "build": "space_runtime",
    "diagnose": "repair",
    "diagnosis": "repair",
    "patch": "repair",
}

DONE_STATUSES = {
    "success",
    "done",
    "completed",
    "passed",
    "full_inference_success",
    "full_inference_candidate_health_passed",
    "manual_hardware_required",
    "technical_blocker",
    "health_only",
}
RUNNING_STATUSES = {"started", "running", "waiting", "pending", "scheduled"}
FAILED_STATUSES = {"failed", "error", "failure"}
CANCELLED_STATUSES = {"cancelled", "canceled"}


def _parse_ts(ts: str | None) -> datetime | None:
    if not ts:
        return None
    try:
        return datetime.fromisoformat(ts.replace("Z", "+00:00"))
    except Exception:
        return None


def _canonical_step(step: str | None) -> str | None:
    if not step:
        return None
    value = str(step)
    return STEP_ALIASES.get(value, value if value in STEP_ORDER else None)


def _is_terminal_status(status: str | None) -> bool:
    s = str(status or "").lower()
    return s in DONE_STATUSES or s in FAILED_STATUSES or s in CANCELLED_STATUSES


def _progress_for_index(index: int) -> int:
    if index <= 0:
        return 3
    if len(STEP_ORDER) <= 1:
        return 0
    return int(round((index / (len(STEP_ORDER) - 1)) * 100))


def progress_from_events(events: list[dict[str, Any]] | None, *, state: dict[str, Any] | None = None) -> dict[str, Any]:
    """Build a monotonic live progress model from worker events + Job state.

    The UI was confusing when late/partial bucket events made steps appear to
    move backwards. This model treats the furthest observed real worker step as
    the source of truth, marks all earlier steps done, keeps only that step
    running while the Job is active, and freezes it as stopped on terminal
    non-success statuses.
    """
    events = events or []
    state = state or {}
    step_status = {step: "pending" for step in STEP_ORDER}
    last_event: dict[str, Any] | None = None
    terminal_status = None
    first_ts = None
    last_ts = None
    furthest_index = -1
    current_step = "bootstrap"

    for event in events:
        if not isinstance(event, dict):
            continue
        step = _canonical_step(event.get("step"))
        status = str(event.get("status") or "").lower()
        ts = _parse_ts(event.get("ts"))
        if ts and not first_ts:
            first_ts = ts
        if ts:
            last_ts = ts
        if not step:
            continue

        last_event = event
        idx = STEP_ORDER.index(step)
        if idx >= furthest_index:
            furthest_index = idx
            current_step = step

        if status in FAILED_STATUSES:
            terminal_status = "failed"
        elif status in CANCELLED_STATUSES:
            terminal_status = "cancelled"
        elif step in {"done", "failure"} and status:
            terminal_status = status

    status_from_state = str(state.get("status") or state.get("job_stage") or state.get("gate_status") or "").lower()
    if status_from_state in CANCELLED_STATUSES:
        terminal_status = "cancelled"
    elif status_from_state in FAILED_STATUSES:
        terminal_status = "failed"
    elif status_from_state in DONE_STATUSES:
        terminal_status = status_from_state

    overall_status = terminal_status or status_from_state or ("running" if events else "not_started")
    terminal = _is_terminal_status(overall_status)

    if furthest_index < 0:
        furthest_index = 0

    if overall_status in DONE_STATUSES or current_step == "done":
        current_step = "done"
        furthest_index = STEP_ORDER.index("done")
        for step in STEP_ORDER:
            if step != "failure":
                step_status[step] = "done"
    else:
        for idx, step in enumerate(STEP_ORDER):
            if idx < furthest_index:
                step_status[step] = "done"
            elif idx == furthest_index:
                if overall_status in FAILED_STATUSES:
                    step_status[step] = "failed"
                elif overall_status in CANCELLED_STATUSES or (terminal and overall_status not in DONE_STATUSES):
                    step_status[step] = "stopped"
                elif overall_status in {"not_started", "idle"}:
                    step_status[step] = "pending"
                else:
                    step_status[step] = "running"
            else:
                step_status[step] = "pending"

    if overall_status in FAILED_STATUSES:
        progress = max(_progress_for_index(furthest_index), 3)
    elif overall_status in CANCELLED_STATUSES:
        progress = max(_progress_for_index(furthest_index), 3)
    elif overall_status in DONE_STATUSES or current_step == "done":
        progress = 100
    else:
        progress = max(3, min(99, _progress_for_index(furthest_index)))

    now = datetime.now(timezone.utc)
    started = first_ts or _parse_ts(str(state.get("started_at") or state.get("created_at") or ""))
    if started:
        end = (last_ts if terminal and last_ts else now)
        elapsed = int((end - started).total_seconds())
    else:
        elapsed = 0

    timeline = [
        {"step": step, "label": STEP_LABELS[step], "status": step_status[step]}
        for step in STEP_ORDER
    ]

    last_message = None
    if last_event:
        last_message = last_event.get("message") or last_event.get("step")

    return {
        "status": overall_status,
        "progress": int(max(0, min(100, progress))),
        "current_step": current_step,
        "current_step_label": STEP_LABELS.get(current_step, current_step),
        "last_event": last_message or "No events yet",
        "elapsed_seconds": max(0, elapsed),
        "eta_seconds": None,
        "timeline": timeline,
        "last_event_raw": last_event or {},
    }