src/timeline_model.py

from __future__ import annotations

from typing import Any

from .progress import progress_from_events

PHASE_ORDER = [
    "start",
    "model",
    "agent",
    "hardware",
    "deploy",
    "live_validation",
    "recovery",
    "archive",
    "done",
]

PHASE_LABELS = {
    "start": "Start",
    "model": "Model",
    "agent": "Agent",
    "hardware": "Hardware",
    "deploy": "Deploy",
    "live_validation": "Live test",
    "recovery": "Recovery",
    "archive": "Archive",
    "done": "Done",
}

PHASE_STEPS = {
    "start": {"token_context", "bootstrap", "dependencies", "auth", "workspace", "node", "pi_install"},
    "model": {"model_analysis", "model_prescan"},
    "agent": {"pi_config", "pi_run", "traces", "pi_model_resolution", "pi_verification"},
    "hardware": {"hardware_strategy", "create_space_hardware", "create_space"},
    "deploy": {"metadata_sanitize", "requirements_sanitize", "upload_files", "space_logs", "space_runtime", "live_wait"},
    "live_validation": {"api_validation", "generation_smoke", "inference_gate"},
    "recovery": {
        "failure_detected",
        "failure_diagnosis",
        "pi_diagnosis",
        "repair_decision",
        "wait_for_logs",
        "factory_rebuild",
        "repair",
        "repair_diagnosis",
        "repair_brief",
        "repair_plan",
        "repair_patch",
        "repair_upload",
        "repair_validation",
    },
    "archive": {"report_write", "artifact_manifest", "anonymous_eval", "eval_publish", "eval_archive"},
    "done": {"done", "failure", "technical_blocker", "manual_hardware_required"},
}

STEP_TO_PHASE = {step: phase for phase, steps in PHASE_STEPS.items() for step in steps}

SUCCESS_STATUSES = {"success", "done", "passed", "completed", "full_inference_success"}
PARTIAL_STATUSES = {"full_inference_candidate_health_passed", "health_only", "partial", "partial_validation", "completed_with_warnings"}
FAILED_STATUSES = {"failed", "failure", "error"}
RUNNING_STATUSES = {"started", "running", "waiting", "pending", "scheduled"}
CANCELLED_STATUSES = {"cancelled", "canceled"}
MANUAL_STATUSES = {"manual_hardware_required", "manual_action_required", "generated_needs_manual_hardware", "waiting_manual_hardware"}
BLOCKED_STATUSES = {"technical_blocker", "blocked"}


def _lower(value: Any) -> str:
    return str(value or "").strip().lower()


def _events(bundle: dict[str, Any]) -> list[dict[str, Any]]:
    return [event for event in (bundle.get("events") or []) if isinstance(event, dict)]


def _event_steps(bundle: dict[str, Any]) -> set[str]:
    return {_lower(event.get("step")) for event in _events(bundle) if event.get("step")}


def _events_for_phase(bundle: dict[str, Any], phase: str) -> list[dict[str, Any]]:
    steps = PHASE_STEPS.get(phase, set())
    return [event for event in _events(bundle) if _lower(event.get("step")) in steps]


def _latest_event_for_phase(bundle: dict[str, Any], phase: str) -> dict[str, Any] | None:
    events = _events_for_phase(bundle, phase)
    return events[-1] if events else None


def _latest_events_by_step(bundle: dict[str, Any], phase: str) -> dict[str, dict[str, Any]]:
    latest: dict[str, dict[str, Any]] = {}
    for event in _events_for_phase(bundle, phase):
        step = _lower(event.get("step"))
        if step:
            latest[step] = event
    return latest


def _has_downstream_proof(bundle: dict[str, Any], phase: str) -> bool:
    steps = _event_steps(bundle)
    phase_index = PHASE_ORDER.index(phase)
    return any(steps.intersection(PHASE_STEPS[p]) for p in PHASE_ORDER[phase_index + 1 :])


def _has_successful_phase_event(bundle: dict[str, Any], phase: str) -> bool:
    return any(_lower(e.get("status")) in SUCCESS_STATUSES for e in _events_for_phase(bundle, phase))


def _has_status(bundle: dict[str, Any], statuses: set[str]) -> bool:
    state = bundle.get("state") or {}
    summary = bundle.get("summary") or {}
    gate = bundle.get("inference_gate") or {}
    smoke = bundle.get("generation_smoke") or {}
    values = {
        _lower(state.get("status")),
        _lower(state.get("gate_status")),
        _lower(summary.get("status")),
        _lower(gate.get("status")),
        _lower(smoke.get("status")),
    }
    values.update(_lower(event.get("status")) for event in _events(bundle))
    return bool(values.intersection(statuses))


def _gate(bundle: dict[str, Any]) -> dict[str, Any]:
    return bundle.get("inference_gate") or {}


def _signals(bundle: dict[str, Any]) -> dict[str, Any]:
    gate = _gate(bundle)
    return gate.get("implementation_signals") if isinstance(gate.get("implementation_signals"), dict) else {}


def _smoke(bundle: dict[str, Any]) -> dict[str, Any]:
    smoke = bundle.get("generation_smoke") or {}
    if smoke:
        return smoke
    gate_smoke = (_gate(bundle).get("generation_smoke") or {})
    return gate_smoke if isinstance(gate_smoke, dict) else {}


def _eval_publish(bundle: dict[str, Any]) -> dict[str, Any]:
    status = bundle.get("eval_publish_status") or bundle.get("eval_publish") or {}
    return status if isinstance(status, dict) else {}


def _build_error_observation(bundle: dict[str, Any]) -> dict[str, Any]:
    observation = bundle.get("build_error_observation") or {}
    return observation if isinstance(observation, dict) else {}


def _final_blocker_label(bundle: dict[str, Any]) -> str:
    observation = _build_error_observation(bundle)
    reason = _lower(observation.get("reason"))
    tail = str(observation.get("tail") or observation.get("first_error") or "")
    if reason or "pyenv install 3.1" in tail or "BUILD FAILED" in tail:
        if "pyenv install 3.1" in tail or "python3.1" in tail.replace(" ", ""):
            return "Build failed: invalid Python version requested pyenv install 3.1"
        return "Build failed after repair"
    blockage = bundle.get("blockage") or {}
    if isinstance(blockage, dict):
        status = str(blockage.get("status") or "").replace("_", " ").strip()
        decision = blockage.get("decision") if isinstance(blockage.get("decision"), dict) else {}
        reason_text = str(decision.get("reason") or blockage.get("reason") or "").strip()
        if reason_text:
            return reason_text[:220]
        if status:
            return status.title()
    return ""


def _is_record_not_ready(status: dict[str, Any]) -> bool:
    return _lower(status.get("reason")) == "record_not_ready"


def _has_terminal_event(bundle: dict[str, Any]) -> bool:
    return any(_lower(event.get("step")) in {"done", "failure", "technical_blocker", "manual_hardware_required"} for event in _events(bundle))


def _latest_recovery_event(bundle: dict[str, Any]) -> dict[str, Any] | None:
    events = _events_for_phase(bundle, "recovery")
    return events[-1] if events else None


def _phase_status_from_events(bundle: dict[str, Any], phase: str) -> str:
    latest_by_step = _latest_events_by_step(bundle, phase)
    if not latest_by_step:
        return "pending"
    statuses = {_lower(event.get("status")) for event in latest_by_step.values()}
    if statuses.intersection(FAILED_STATUSES):
        return "failed"
    if statuses.intersection({"warning"}) or statuses.intersection(PARTIAL_STATUSES):
        return "warning"
    if statuses.intersection(RUNNING_STATUSES):
        # A historical started/waiting event should not keep a terminal phase
        # running when another step in that phase already succeeded. The live
        # Job logs often contain bootstrap/api_validation started events without
        # a matching success event for the same step.
        if statuses.intersection(SUCCESS_STATUSES):
            return "complete"
        return "running"
    if statuses.intersection(SUCCESS_STATUSES):
        return "complete"
    return "complete"


def _verdict(bundle: dict[str, Any]) -> str:
    eval_record = bundle.get("eval_record") or {}
    if isinstance(eval_record, dict) and eval_record.get("verdict"):
        return _lower(eval_record.get("verdict"))
    gate = _gate(bundle)
    state = bundle.get("state") or {}
    for value in (gate.get("status"), state.get("gate_status"), state.get("status"), (bundle.get("summary") or {}).get("status")):
        normalized = _lower(value)
        if normalized:
            if normalized in SUCCESS_STATUSES:
                return "full_inference_success" if normalized == "full_inference_success" else "success"
            if normalized in PARTIAL_STATUSES:
                return "partial_validation"
            if normalized in MANUAL_STATUSES:
                return "manual_action_required"
            if normalized in BLOCKED_STATUSES:
                return "technical_blocker"
            if normalized in FAILED_STATUSES:
                return "failed"
            if normalized in CANCELLED_STATUSES:
                return "cancelled"
            return normalized
    return "running" if _events(bundle) else "unknown"


def _final_visual_status(verdict: str) -> str:
    if verdict in {"full_inference_success", "success", "passed"}:
        return "success"
    if verdict in {"partial_validation", "partial", "health_only", "completed_with_warnings"}:
        return "warn"
    if verdict in {"failed", "failure", "technical_blocker"}:
        return "error"
    if verdict in {"manual_action_required", "cancelled"}:
        return "stopped"
    return "running"


def _status_label(verdict: str) -> tuple[str, str]:
    if verdict == "full_inference_success":
        return "Full inference success", "Space boots, generation passed, and latency was measured."
    if verdict == "partial_validation":
        return "Completed with partial validation", "The run finished, but full generation was not verified."
    if verdict == "technical_blocker":
        return "Technical blocker", "The run found a technical blocker."
    if verdict == "manual_action_required":
        return "Manual action required", "The run needs user action before validation can continue."
    if verdict == "failed":
        return "Failed", "The run did not complete successfully."
    if verdict == "cancelled":
        return "Cancelled", "The run was cancelled."
    if verdict in {"running", "pending", "unknown"}:
        return "Running", "The pipeline is still in progress."
    return verdict.replace("_", " ").title(), ""


def _hardware_warning(bundle: dict[str, Any]) -> dict[str, Any] | None:
    events = _events_for_phase(bundle, "hardware")
    preferred_failed = any(_lower(e.get("step")) == "create_space_hardware" and _lower(e.get("status")) in FAILED_STATUSES for e in events)
    fallback_ok = any(_lower(e.get("step")) in {"create_space", "create_space_hardware"} and _lower(e.get("status")) in SUCCESS_STATUSES and _lower((e.get("data") or {}).get("hardware")) for e in events)
    if preferred_failed and fallback_ok:
        return {"code": "hardware_fallback_used", "label": "Fallback GPU used", "detail": "Preferred hardware was unavailable; the factory continued on fallback hardware."}
    return None


def _pi_model_warning(bundle: dict[str, Any]) -> dict[str, Any] | None:
    for event in _events(bundle):
        if _lower(event.get("step")) == "pi_model_resolution" and _lower(event.get("status")) == "warning":
            data = event.get("data") or {}
            requested = data.get("requested_model") or data.get("configured_model") or "requested model"
            effective = data.get("effective_model") or "observed model"
            return {"code": "pi_model_changed", "label": "Pi assistant model changed", "detail": f"{requested} → {effective}"}
    return None


def _logs_warning(bundle: dict[str, Any]) -> dict[str, Any] | None:
    logs = bundle.get("space_logs_index") or bundle.get("space_logs") or {}
    if not isinstance(logs, dict):
        return None
    quality = _lower(logs.get("log_quality") or logs.get("quality"))
    if quality and quality not in {"full", "complete"}:
        return {"code": "logs_partial", "label": "Space logs partial", "detail": f"Log quality: {quality}."}
    return None


def _archive_warning(bundle: dict[str, Any], *, terminal: bool) -> dict[str, Any] | None:
    status = _eval_publish(bundle)
    if not status or status.get("published") is True:
        return None
    # record_not_ready is the expected state while a run is still executing.
    # It should not become a global warning until the run has actually ended.
    if not terminal and _is_record_not_ready(status):
        return None
    if status.get("attempted"):
        return {"code": "eval_archive_not_published", "label": "Eval archive not published", "detail": str(status.get("reason") or "Backend archive copy did not complete.")}
    if terminal:
        return {"code": "eval_archive_pending", "label": "Eval archive pending", "detail": "Backend archive copy has not been confirmed yet."}
    return None


def _build_phase_details(bundle: dict[str, Any], phase: str) -> list[dict[str, Any]]:
    details: list[dict[str, Any]] = []
    if phase == "live_validation":
        signals = _signals(bundle)
        smoke = _smoke(bundle)
        if signals.get("health_passed") is True:
            details.append({"label": "Health passed", "status": "complete"})
        elif signals.get("health_passed") is False:
            details.append({"label": "Health not confirmed", "status": "warning"})
        if signals.get("generation_smoke_passed") is True or _lower(smoke.get("status")) == "success":
            latency = smoke.get("latency_seconds") or smoke.get("observed_latency_seconds")
            text = f"Generation passed in {latency}s" if latency is not None else "Generation passed"
            details.append({"label": text, "status": "complete"})
        elif smoke:
            reason = smoke.get("failure_type") or smoke.get("error") or "Generation was not verified"
            details.append({"label": f"Generation not verified: {reason}", "status": "warning"})
        blocker = _final_blocker_label(bundle)
        if blocker:
            details.append({"label": blocker, "status": "failed"})
        recommendation = smoke.get("recommended_zero_gpu_duration_seconds") or (_gate(bundle).get("zero_gpu_duration_recommendation") or {}).get("recommended_zero_gpu_duration_seconds")
        if recommendation:
            details.append({"label": f"ZeroGPU duration recommendation: {recommendation}s", "status": "complete"})
    elif phase == "archive":
        eval_status = _eval_publish(bundle)
        if eval_status:
            if eval_status.get("published") is True:
                details.append({"label": "Eval archive published", "status": "complete"})
                if eval_status.get("archive_relative_path"):
                    details.append({"label": str(eval_status.get("archive_relative_path")), "status": "info"})
            elif eval_status.get("attempted"):
                if _is_record_not_ready(eval_status):
                    details.append({"label": "Eval archive pending until the run completes", "status": "pending"})
                else:
                    details.append({"label": f"Eval archive not published: {eval_status.get('reason') or 'unknown'}", "status": "warning"})
            else:
                details.append({"label": "Eval archive pending", "status": "pending"})
    elif phase == "recovery":
        steps = _event_steps(bundle)
        if not steps.intersection(PHASE_STEPS["recovery"]):
            details.append({"label": "Recovery not needed", "status": "skipped"})
        else:
            failed_steps = {
                _lower(e.get("step"))
                for e in _events(bundle)
                if _lower(e.get("step")) in PHASE_STEPS["recovery"] and _lower(e.get("status")) in FAILED_STATUSES
            }
            for label, candidates in [
                ("Diagnose", {"failure_diagnosis", "pi_diagnosis"}),
                ("Decide", {"repair_decision"}),
                ("Patch", {"repair", "repair_patch", "repair_upload"}),
                ("Revalidate", {"repair_validation"}),
            ]:
                status = "failed" if failed_steps.intersection(candidates) else "complete" if steps.intersection(candidates) else "pending"
                details.append({"label": label, "status": status})
            blocker = _final_blocker_label(bundle)
            if blocker:
                details.append({"label": blocker, "status": "failed"})
    else:
        latest = _latest_event_for_phase(bundle, phase)
        if latest:
            details.append({"label": str(latest.get("message") or latest.get("step") or PHASE_LABELS[phase]), "status": _lower(latest.get("status")) or "info"})
    return details


def _phase_summary(bundle: dict[str, Any], phase: str, status: str) -> str:
    if phase == "agent":
        warning = _pi_model_warning(bundle)
        if warning:
            return warning["label"]
    if phase == "hardware":
        warning = _hardware_warning(bundle)
        if warning:
            return warning["label"]
    if phase == "live_validation":
        blocker = _final_blocker_label(bundle)
        if blocker and _final_visual_status(_verdict(bundle)) == "error":
            return "Live validation incomplete"
        signals = _signals(bundle)
        smoke = _smoke(bundle)
        if signals.get("generation_smoke_passed") is True or _lower(smoke.get("status")) == "success":
            latency = smoke.get("latency_seconds") or smoke.get("observed_latency_seconds")
            return f"Generation passed in {latency}s" if latency is not None else "Generation passed"
        if signals.get("health_passed") is True:
            return "Health passed; generation not verified"
        return "Live validation pending" if status in {"pending", "running"} else "Live validation incomplete"
    if phase == "recovery":
        steps = _event_steps(bundle)
        if not steps.intersection(PHASE_STEPS["recovery"]):
            return "Not needed"
        if any(_lower(e.get("step")) == "repair_validation" and _lower(e.get("status")) in FAILED_STATUSES for e in _events(bundle)):
            return "Patch attempted; final blocker remains"
        if any(_lower(e.get("step")) == "repair_validation" and _lower(e.get("status")) in SUCCESS_STATUSES for e in _events(bundle)):
            return "Repair revalidated"
        latest = _latest_recovery_event(bundle) or {}
        latest_step = _lower(latest.get("step"))
        if latest_step in {"repair_validation"}:
            return "Revalidating repair"
        if latest_step in {"repair", "repair_plan", "repair_patch", "repair_upload"}:
            return "Repairing"
        return "Diagnosing repair"
    if phase == "archive":
        eval_status = _eval_publish(bundle)
        if eval_status.get("published") is True:
            return "Eval archive published"
        if eval_status.get("attempted") and not _is_record_not_ready(eval_status):
            return "Eval archive attempted"
        return "Archive pending" if status in {"pending", "running"} else "Final artifacts written"
    if phase == "done":
        label, _ = _status_label(_verdict(bundle))
        return label
    latest = _latest_event_for_phase(bundle, phase)
    if latest:
        return str(latest.get("message") or PHASE_LABELS[phase])
    if status == "skipped":
        return "Not needed"
    return "Pending" if status == "pending" else PHASE_LABELS[phase]


def _phase_status(bundle: dict[str, Any], phase: str, verdict: str) -> str:
    steps = _event_steps(bundle)
    terminal = _final_visual_status(verdict) in {"success", "warn", "error", "stopped"}

    if phase == "done":
        if verdict in {"full_inference_success", "success", "passed"}:
            return "complete"
        if verdict in {"partial_validation", "partial", "completed_with_warnings", "health_only"}:
            return "warning"
        if verdict in {"failed", "technical_blocker"}:
            return "failed"
        if verdict in {"manual_action_required", "cancelled"}:
            return "stopped"
        return "pending"

    if phase == "recovery":
        if not steps.intersection(PHASE_STEPS["recovery"]):
            return "skipped" if terminal else "pending"
        if any(_lower(e.get("step")) == "repair_validation" and _lower(e.get("status")) in FAILED_STATUSES for e in _events(bundle)):
            return "failed"
        if any(_lower(e.get("step")) == "repair_validation" and _lower(e.get("status")) in SUCCESS_STATUSES for e in _events(bundle)):
            return "complete"
        latest_recovery = _latest_recovery_event(bundle) or {}
        latest_status = _lower(latest_recovery.get("status"))
        if latest_status in FAILED_STATUSES:
            return "failed"
        return "warning" if terminal else "running"

    if phase == "live_validation":
        blocker = _final_blocker_label(bundle)
        if blocker and _final_visual_status(_verdict(bundle)) == "error":
            return "Live validation incomplete"
        signals = _signals(bundle)
        smoke = _smoke(bundle)
        if signals.get("generation_smoke_passed") is True or _lower(smoke.get("status")) == "success":
            return "complete"
        if signals.get("health_passed") is True or smoke:
            return "warning"
        return _phase_status_from_events(bundle, phase)

    if phase == "archive":
        eval_status = _eval_publish(bundle)
        if eval_status.get("published") is True:
            return "complete"
        if eval_status.get("attempted"):
            if not terminal and _is_record_not_ready(eval_status):
                return "pending"
            return "warning"
        if terminal:
            return "warning" if bundle.get("eval_record") else "pending"
        return "pending"

    if phase == "agent" and _pi_model_warning(bundle):
        return "warning"
    if phase == "hardware" and _hardware_warning(bundle):
        return "warning"

    event_status = _phase_status_from_events(bundle, phase)
    if event_status != "pending":
        if event_status == "running" and terminal and (_has_downstream_proof(bundle, phase) or _has_successful_phase_event(bundle, phase)):
            return "complete"
        return event_status

    # On terminal runs, phases before the first missing/optional later phase can
    # be considered complete when downstream proofs exist.
    if terminal and _has_downstream_proof(bundle, phase):
        return "complete"
    return "pending"


def build_run_timeline_model(bundle: dict[str, Any]) -> dict[str, Any]:
    """Return the canonical product progress/timeline projection for a run.

    This model is deliberately derived from persisted run artifacts and events,
    not from transient DOM state. It separates process completion from the final
    result verdict so a completed run can be success, partial, failed, blocked,
    or manual-action-required without the progress bar lying.
    """
    bundle = bundle or {}
    state = bundle.get("state") or {}
    progress = progress_from_events(_events(bundle), state=state)
    verdict = _verdict(bundle)
    visual_status = _final_visual_status(verdict)
    label, subtitle = _status_label(verdict)
    terminal = visual_status in {"success", "warn", "error", "stopped"}
    percent = 100 if terminal else int(progress.get("progress") or 0)

    warnings = [w for w in (_pi_model_warning(bundle), _hardware_warning(bundle), _logs_warning(bundle), _archive_warning(bundle, terminal=terminal)) if w]
    phases: list[dict[str, Any]] = []
    for phase in PHASE_ORDER:
        status = _phase_status(bundle, phase, verdict)
        phases.append(
            {
                "id": phase,
                "label": PHASE_LABELS[phase],
                "status": status,
                "summary": _phase_summary(bundle, phase, status),
                "details": _build_phase_details(bundle, phase),
            }
        )

    return {
        "schema_version": "run_timeline_model.v1",
        "progress": {
            "percent": max(0, min(100, percent)),
            "visual_status": visual_status,
            "label": label,
            "subtitle": subtitle,
            "process_status": progress.get("status") or "unknown",
            "verdict": verdict,
            "terminal": terminal,
        },
        "phases": phases,
        "warnings": warnings,
    }