"""Private matplotlib animation queue helpers."""

from typing import TYPE_CHECKING, Any

import numpy as np
from nxslib.nxscope import DNxscopeStreamBlock

from nxscli_mpl._plot_constants import ANIMATION_FRAME_DRAIN_LIMIT

if TYPE_CHECKING:
    from nxscli.idata import PluginQueueData


def fetch_animation_frame(
    qdata: "PluginQueueData",
    *,
    count: int,
    limit: int = ANIMATION_FRAME_DRAIN_LIMIT,
    stop_on_trigger: bool = False,
) -> tuple[
    list[np.ndarray[Any, Any]],
    list[np.ndarray[Any, Any]],
    int,
    float | None,
]:
    """Drain one animation frame from the queue and return updated counters."""
    x_chunks: list[np.ndarray[Any, Any]] = []
    y_chunks: list[list[np.ndarray[Any, Any]]] = [
        [] for _ in range(qdata.vdim)
    ]
    next_count = count
    trigger_x: float | None = None

    for _ in range(limit):
        data = qdata.queue_get(block=False)
        if not data:
            break
        if not isinstance(data, list):
            raise RuntimeError("plot animation queue payload must be list")
        event_getter = getattr(qdata, "pop_trigger_event", None)
        event = event_getter() if callable(event_getter) else None
        batch_start = next_count
        if event is not None and trigger_x is None:
            trigger_x = batch_start + event.sample_index
        had_samples = False

        for block in data:
            if not isinstance(block, DNxscopeStreamBlock):
                raise RuntimeError(
                    "plot animation requires DNxscopeStreamBlock payload"
                )
            block_data = block.data
            assert isinstance(block_data, np.ndarray)
            nsamples = int(block_data.shape[0])
            if nsamples == 0:
                continue
            had_samples = True
            xr = np.arange(next_count, next_count + nsamples)
            x_chunks.append(xr)
            for i in range(qdata.vdim):
                y_chunks[i].append(block_data[:, i])
            next_count += nsamples
        if stop_on_trigger and event is not None and had_samples:
            break

    xcat = (
        np.concatenate(x_chunks)
        if x_chunks
        else np.empty((0,), dtype=np.int64)
    )
    xdata = [xcat for _ in range(qdata.vdim)]
    ydata = [
        (
            np.concatenate(chunks)
            if chunks
            else np.empty((0,), dtype=np.float64)
        )
        for chunks in y_chunks
    ]
    return xdata, ydata, next_count, trigger_x
