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"
    current_status = ""
    current_event: dict[str, Any] | None = None
    latest_signal_step = ""
    latest_signal_status = ""
    latest_signal_event: dict[str, Any] | None = None
    latest_signal_ts: datetime | None = None

    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)
        # Never let older / lower-index failed events poison the whole run once
        # later validation or repair steps have appeared.
        if idx >= furthest_index:
            furthest_index = idx
            current_step = step
            current_status = status
            current_event = event
        if ts is None or latest_signal_ts is None or ts >= latest_signal_ts:
            latest_signal_ts = ts
            latest_signal_step = step
            latest_signal_status = status
            latest_signal_event = event

    if latest_signal_status in RUNNING_STATUSES:
        current_step = latest_signal_step or current_step
        current_status = latest_signal_status
        current_event = latest_signal_event or current_event
        if current_step in STEP_ORDER:
            furthest_index = STEP_ORDER.index(current_step)
        terminal_status = None

    if current_step == "done" and current_status:
        terminal_status = current_status
    elif current_step == "failure" and current_status:
        terminal_status = "failed" if current_status in FAILED_STATUSES or current_status in DONE_STATUSES else current_status
    elif current_status in CANCELLED_STATUSES:
        terminal_status = "cancelled"
    elif current_status in FAILED_STATUSES:
        terminal_status = "failed"

    status_from_state = str(state.get("status") or state.get("gate_status") or "").lower()
    job_stage = str(state.get("job_stage") or "").lower()
    if status_from_state in CANCELLED_STATUSES:
        terminal_status = "cancelled"
    elif status_from_state in DONE_STATUSES:
        terminal_status = status_from_state
    elif status_from_state in FAILED_STATUSES and (not events or current_status in FAILED_STATUSES or current_step == "failure"):
        terminal_status = "failed"
    elif job_stage in CANCELLED_STATUSES:
        terminal_status = "cancelled"
    elif job_stage in FAILED_STATUSES:
        # HF Job failed is authoritative only when the worker did not emit a
        # later waiting/running validation/manual state. This avoids showing a
        # recovered run as failed, but still turns genuinely stopped jobs red.
        if not events or current_status in FAILED_STATUSES or current_step == "failure":
            terminal_status = "failed"
        elif current_status not in RUNNING_STATUSES:
            terminal_status = "failed"

    if terminal_status:
        overall_status = terminal_status
    elif current_status in RUNNING_STATUSES:
        overall_status = "running"
    elif status_from_state in RUNNING_STATUSES:
        overall_status = "running"
    elif events:
        overall_status = "running"
    else:
        overall_status = "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 current_event:
        last_message = current_event.get("message") or current_event.get("step")
    elif 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",
        "visual_status": "error" if overall_status in FAILED_STATUSES else ("stopped" if overall_status in CANCELLED_STATUSES else ("success" if overall_status in DONE_STATUSES and overall_status not in {"manual_hardware_required", "technical_blocker"} else ("stopped" if overall_status in {"manual_hardware_required", "technical_blocker"} else ("running" if overall_status == "running" else "neutral")))),
        "elapsed_seconds": max(0, elapsed),
        "eta_seconds": None,
        "timeline": timeline,
        "last_event_raw": last_event or {},
    }