from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any

STEP_ORDER = [
    "bucket_ready",
    "job_launched",
    "dependencies",
    "auth",
    "model_analysis",
    "workspace",
    "node",
    "pi_install",
    "pi_config",
    "pi_run",
    "create_space",
    "upload_files",
    "hardware",
    "api_validation",
    "inference_gate",
    "report_write",
    "done",
]

STEP_LABELS = {
    "bucket_ready": "Bucket ready",
    "job_launched": "Job launched",
    "dependencies": "Dependencies",
    "auth": "Authenticated",
    "model_analysis": "Model analysis",
    "workspace": "Workspace",
    "node": "Node/npm",
    "pi_install": "Pi installed",
    "pi_config": "Pi configured",
    "pi_run": "Pi running",
    "create_space": "Create Space",
    "upload_files": "Upload files",
    "hardware": "Hardware request",
    "api_validation": "API validation",
    "inference_gate": "Gate",
    "report_write": "Report",
    "done": "Done",
}

STEP_PROGRESS = {
    "bucket_ready": 5,
    "job_launched": 8,
    "dependencies": 12,
    "auth": 16,
    "model_analysis": 22,
    "workspace": 28,
    "node": 34,
    "pi_install": 40,
    "pi_config": 44,
    "pi_run": 62,
    "create_space": 72,
    "upload_files": 80,
    "hardware": 86,
    "api_validation": 92,
    "inference_gate": 96,
    "report_write": 98,
    "done": 100,
}

STEP_ALIASES = {
    "bootstrap": "job_launched",
    "dependencies": "dependencies",
    "auth": "auth",
    "model_analysis": "model_analysis",
    "workspace": "workspace",
    "node": "node",
    "pi_install": "pi_install",
    "pi_config": "pi_config",
    "pi_run": "pi_run",
    "create_space": "create_space",
    "upload_files": "upload_files",
    "hardware_preferred": "hardware",
    "hardware_fallback": "hardware",
    "hardware": "hardware",
    "api_validation": "api_validation",
    "inference_gate": "inference_gate",
    "report_write": "report_write",
    "done": "done",
    "failure": "done",
}

DONE_STATUSES = {"success", "done", "completed", "passed", "full_inference_success", "full_inference_candidate_health_passed", "manual_hardware_required", "technical_blocker"}
RUNNING_STATUSES = {"started", "running", "waiting"}
FAILED_STATUSES = {"failed", "error"}


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
    return STEP_ALIASES.get(step, step if step in STEP_PROGRESS else None)


def progress_from_events(events: list[dict[str, Any]] | None, *, state: dict[str, Any] | None = None) -> dict[str, Any]:
    """Build a stable UI progress model from worker events.jsonl + optional state.json."""
    events = events or []
    state = state or {}
    step_status = {step: "pending" for step in STEP_ORDER}
    last_event = None
    current_step = "job_launched"
    terminal_status = None
    first_ts = None
    last_ts = None

    for event in events:
        if not isinstance(event, dict):
            continue
        last_event = event
        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
        current_step = step
        if status in FAILED_STATUSES:
            step_status[step] = "failed"
            terminal_status = "failed"
        elif status in RUNNING_STATUSES:
            if step_status.get(step) != "done":
                step_status[step] = "running"
        elif status in DONE_STATUSES or status:
            step_status[step] = "done"

        if step == "done" and status:
            terminal_status = status

    # Mark all previous steps as done up to the current running/done step.
    current_index = STEP_ORDER.index(current_step) if current_step in STEP_ORDER else 0
    for step in STEP_ORDER[:current_index]:
        if step_status[step] == "pending":
            step_status[step] = "done"

    if terminal_status and terminal_status != "failed":
        current_step = "done"
        for step in STEP_ORDER:
            if step_status[step] != "failed":
                step_status[step] = "done"

    status_from_state = state.get("status") or state.get("gate_status")
    overall_status = terminal_status or status_from_state or ("running" if events else "not_started")

    if overall_status in {"failed", "error"}:
        progress = max(STEP_PROGRESS.get(current_step, 8), 8)
    elif current_step == "done" or overall_status in DONE_STATUSES:
        progress = 100
    else:
        progress = STEP_PROGRESS.get(current_step, 8)
        if step_status.get(current_step) == "running":
            previous = STEP_ORDER[max(0, current_index - 1)] if current_index > 0 else current_step
            progress = max(STEP_PROGRESS.get(previous, 0) + 2, progress - 8)

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

    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 {},
    }