"""Server-Sent Events (SSE) generator for real-time progress streaming to web clients."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/streaming.ipynb.

# %% auto 0
__all__ = ['sse_stream']

# %% ../nbs/streaming.ipynb 3
import json, time
from typing import Iterator
from .progress_monitor import ProgressMonitor

# %% ../nbs/streaming.ipynb 5
def sse_stream(
    monitor: ProgressMonitor,  # Progress monitor instance to read job updates from
    job_id: str,  # Unique identifier of the job to stream
    interval: float = 0.25,  # Polling interval in seconds for checking progress updates
    heartbeat: float = 15.0,  # Seconds between keep-alive messages when no updates
    wait_for_start: bool = True,  # Whether to wait for job to start before ending stream
    start_timeout: float = 5.0,  # Maximum seconds to wait for job to start if wait_for_start is True
) -> Iterator[str]:  # SSE-formatted strings ready to send to client
    """
    Framework-agnostic SSE generator.
    - Yields 'data: {json}\\n\\n' when progress changes.
    - Sends ': keep-alive' comments every `heartbeat` seconds when idle.
    - If `wait_for_start` is True, it will wait up to `start_timeout` for
      the first snapshot before ending (avoids race at job startup).
    """
    last = None
    start_ts = time.time()
    last_hb = start_ts

    while True:
        snap = monitor.snapshot(job_id)
        now = time.time()

        if not snap:
            if wait_for_start and (now - start_ts) < start_timeout:
                # keep the connection alive while waiting for first update
                if now - last_hb >= heartbeat:
                    yield ": waiting\n\n"
                    last_hb = now
                time.sleep(min(interval, 0.1))
                continue
            else:
                yield "event: end\ndata: {}\n\n"
                break

        latest = snap["latest"]
        changed = (
            not last
            or latest.current != getattr(last, "current", None)
            or latest.progress != getattr(last, "progress", None)
        )

        if changed:
            payload = {
                "progress": snap["overall_progress"],
                "completed": snap["completed"],
                "bars": {
                    k: {"desc": v.description, "pct": v.progress, "cur": v.current, "tot": v.total}
                    for k, v in snap["bars"].items()
                },
            }
            yield f"data: {json.dumps(payload)}\n\n"
            last = latest
            last_hb = now
        elif now - last_hb >= heartbeat:
            yield ": keep-alive\n\n"
            last_hb = now

        if snap["completed"]:
            break

        time.sleep(interval)
