"""
MeasureIt backend for sweep operations.

Handles all MeasureIt-related operations including:
- Sweep status monitoring
- Waiting for sweep completion
- Killing sweeps
"""

import asyncio
import time
import logging
from typing import Callable, Coroutine, Dict, Any, Optional

# Async callback: (progress, total, message) -> None
ProgressCallback = Callable[[float, float, str], Coroutine[Any, Any, None]]

from ...backend.base import BaseBackend, SharedState

logger = logging.getLogger(__name__)

# Delay in seconds between checks for wait_for_all_sweeps and wait_for_sweep
WAIT_DELAY = 1.0

# Import MeasureIt types with fallback
try:
    from measureit.sweep.base_sweep import BaseSweep
    from measureit.sweep.progress import SweepState
    from measureit.tools.sweep_queue import SweepQueue

    SWEEP_STATE_ERROR = SweepState.ERROR.value  # "error"
    SWEEP_STATE_RUNNING = (SweepState.RAMPING.value, SweepState.RUNNING.value)
except ImportError:  # pragma: no cover - MeasureIt optional
    BaseSweep = None  # type: ignore[assignment]
    SweepQueue = None  # type: ignore[assignment]
    SWEEP_STATE_ERROR = "error"
    SWEEP_STATE_RUNNING = ("ramping", "running")


def _is_sweep(obj: Any) -> bool:
    """Check if an object is a MeasureIt sweep (BaseSweep or SweepQueue).

    Args:
        obj: Object to check

    Returns:
        True if obj is a BaseSweep or SweepQueue instance
    """
    if BaseSweep is not None and isinstance(obj, BaseSweep):
        return True
    if SweepQueue is not None and isinstance(obj, SweepQueue):
        return True
    return False


def _get_progress_state(sweep: Any) -> Optional[Any]:
    """Get the progress state from a sweep object.

    Handles both BaseSweep (direct progressState attribute) and SweepQueue
    (accesses current_sweep.progressState; note that SweepQueue also has a
    state() method that returns the same value).

    Args:
        sweep: A BaseSweep or SweepQueue instance

    Returns:
        The progressState object, or None if not available
    """
    # BaseSweep has direct progressState attribute
    if BaseSweep is not None and isinstance(sweep, BaseSweep):
        return sweep.progressState

    # SweepQueue delegates to current_sweep's progressState
    if SweepQueue is not None and isinstance(sweep, SweepQueue):
        if hasattr(sweep, "current_sweep") and sweep.current_sweep:
            return sweep.current_sweep.progressState

    return None


def _get_sweep_queue_status(sweep: Any) -> Optional[Dict[str, Any]]:
    """Get comprehensive status from a SweepQueue using its status() method.

    Returns:
        Dict with effective_state, queue_length, last_error, etc., or None if not a SweepQueue
        or if status() is unavailable.
    """
    if SweepQueue is None or not isinstance(sweep, SweepQueue):
        return None

    if hasattr(sweep, "status"):
        return sweep.status()

    return None


class MeasureItBackend(BaseBackend):
    """Backend for MeasureIt sweep operations."""

    def __init__(self, state: SharedState):
        """Initialize MeasureIt backend.

        Args:
            state: SharedState instance containing shared resources
        """
        super().__init__(state)

    async def get_measureit_status(self) -> Dict[str, Any]:
        """Check if any measureit sweep is currently running.

        Returns information about active measureit sweeps in the notebook namespace,
        including sweep type, status, and basic configuration if available.

        Returns:
            Dict containing:
                - active: bool - whether any sweep is active
                - sweeps: Dict mapping variable names to Dicts of active sweep information:
                    "variable_name" (str), "module" (str), "state" (str), "progress" (float or None), "time_elapsed" (float or None), "time_remaining" (float or None)
                - error: str (if any error occurred)
        """
        try:
            if BaseSweep is None:
                return {
                    "active": False,
                    "sweeps": {},
                    "error": "MeasureIt library not available",
                }

            result = {"active": False, "sweeps": {}}

            # Look for MeasureIt sweep objects in the namespace
            for var_name, var_value in self.namespace.items():
                # Skip private/internal variables
                if var_name.startswith("_"):
                    continue

                if _is_sweep(var_value):
                    progress_state = _get_progress_state(var_value)

                    # Skip individual sweeps (BaseSweep) managed by a SweepQueue
                    # Don't skip SweepQueue objects themselves
                    is_base_sweep = BaseSweep is not None and isinstance(
                        var_value, BaseSweep
                    )
                    if is_base_sweep and progress_state:
                        if getattr(progress_state, "is_queued", False):
                            continue

                    sweep_info = {
                        "variable_name": var_name,
                        "module": getattr(var_value, "__module__", ""),
                    }

                    # Check if this is a SweepQueue
                    is_sweep_queue = SweepQueue is not None and isinstance(
                        var_value, SweepQueue
                    )

                    if is_sweep_queue:
                        # Use SweepQueue.status() for comprehensive state
                        queue_status = _get_sweep_queue_status(var_value)
                        if queue_status:
                            sweep_info["type"] = "SweepQueue"
                            sweep_info["state"] = queue_status.get(
                                "effective_state", "idle"
                            )
                            sweep_info["queue_length"] = queue_status.get(
                                "queue_length", 0
                            )
                            sweep_info["current_sweep_state"] = queue_status.get(
                                "current_sweep_state"
                            )
                            if queue_status.get("last_error"):
                                sweep_info["error_message"] = queue_status["last_error"]

                            # Add progress from current sweep if available
                            if progress_state:
                                sweep_info["progress"] = progress_state.progress
                                sweep_info["time_elapsed"] = progress_state.time_elapsed
                                sweep_info["time_remaining"] = (
                                    progress_state.time_remaining
                                )

                            queue_active = sweep_info["state"] in (
                                "running",
                                "pending",
                                "paused",
                            )
                            result["active"] = result["active"] or queue_active
                            result["sweeps"][var_name] = sweep_info
                        elif progress_state:
                            # Fallback for older MeasureIt without status()
                            sweep_info["type"] = "SweepQueue"
                            sweep_info["state"] = progress_state.state.value
                            sweep_info["progress"] = progress_state.progress
                            sweep_info["time_elapsed"] = progress_state.time_elapsed
                            sweep_info["time_remaining"] = progress_state.time_remaining

                            # In fallback mode, treat running states as active
                            # Note: can't detect "pending" without status(), so may
                            # misclassify idle queue with pending items as inactive
                            result["active"] = result["active"] or (
                                sweep_info["state"] in SWEEP_STATE_RUNNING
                            )

                            error_message = getattr(
                                progress_state, "error_message", None
                            )
                            if error_message:
                                sweep_info["error_message"] = error_message

                            result["sweeps"][var_name] = sweep_info
                    elif progress_state:
                        # Regular BaseSweep
                        sweep_info["state"] = progress_state.state.value
                        sweep_info["progress"] = progress_state.progress
                        sweep_info["time_elapsed"] = progress_state.time_elapsed
                        sweep_info["time_remaining"] = progress_state.time_remaining

                        result["active"] = result["active"] or (
                            sweep_info["state"] in SWEEP_STATE_RUNNING
                        )

                        # Capture error information if present
                        error_message = getattr(progress_state, "error_message", None)
                        if error_message:
                            sweep_info["error_message"] = error_message

                        result["sweeps"][var_name] = sweep_info

            return result

        except Exception as e:
            logger.error(f"Error checking MeasureIt status: {e}")
            return {"active": False, "sweeps": {}, "error": str(e)}

    async def kill_sweep(self, var_name: str) -> Dict[str, Any]:
        """Kill a running MeasureIt sweep to release resources.

        UNSAFE: This tool stops a running sweep, which may leave instruments
        in an intermediate state. Use when a sweep needs to be terminated
        due to timeout, error, or user request.

        Args:
            var_name: Name of the sweep variable in the notebook namespace

        Returns:
            Dict containing:
                success: bool - whether the kill was successful (sweep stopped)
                sweep_name: str - name of the sweep
                previous_state: str - state before kill
                new_state: str - state after kill attempt
                error: str (if kill failed or sweep still running)
        """
        try:
            if BaseSweep is None:
                return {
                    "success": False,
                    "sweep_name": var_name,
                    "error": "MeasureIt library not available",
                }

            # Check if the variable exists in namespace
            if var_name not in self.namespace:
                return {
                    "success": False,
                    "sweep_name": var_name,
                    "error": f"Variable '{var_name}' not found in namespace",
                }

            sweep = self.namespace[var_name]

            # Verify it's a sweep object
            if not _is_sweep(sweep):
                return {
                    "success": False,
                    "sweep_name": var_name,
                    "error": f"Variable '{var_name}' is not a MeasureIt sweep (got {type(sweep).__name__})",
                }

            # Check if this is a SweepQueue
            is_sweep_queue = SweepQueue is not None and isinstance(sweep, SweepQueue)

            # Get progress state (handles both BaseSweep and SweepQueue)
            progress_state = _get_progress_state(sweep)

            # For SweepQueue without current_sweep, use status() to get state
            if not progress_state:
                if is_sweep_queue:
                    queue_status = _get_sweep_queue_status(sweep)
                    if queue_status:
                        previous_state = queue_status.get("effective_state", "idle")
                    else:
                        previous_state = "idle"
                else:
                    return {
                        "success": False,
                        "sweep_name": var_name,
                        "error": f"Could not access progress state for '{var_name}' (type: {type(sweep).__name__})",
                    }
            else:
                # Get state before kill
                previous_state = progress_state.state.value

            # Execute sweep.kill() via Qt's thread-safe mechanism
            #
            # The MCP server runs in a separate thread from the sweep's Qt thread.
            # We use QMetaObject.invokeMethod with QueuedConnection to queue the
            # kill directly on the sweep's Qt thread - the same mechanism MeasureIt
            # uses internally. This bypasses Tornado/asyncio entirely.
            #
            # IMPORTANT: Always use this Qt proxy mechanism, even for stopped sweeps,
            # because calling sweep.kill() directly from the MCP thread can cause
            # deadlocks or crashes due to Qt threading requirements.
            try:
                from PyQt5.QtCore import QObject, QMetaObject, Qt, pyqtSlot

                QueuedConnection = Qt.QueuedConnection
            except ImportError:
                from PyQt6.QtCore import QObject, QMetaObject, Qt, pyqtSlot

                # PyQt6 moved ConnectionType to an enum
                QueuedConnection = Qt.ConnectionType.QueuedConnection

            # Create or reuse a kill proxy that lives on the sweep's thread
            proxy = getattr(sweep, "_mcp_kill_proxy", None)
            if proxy is None:

                class _KillProxy(QObject):
                    def __init__(self, target_sweep, use_kill_all=False):
                        super().__init__()
                        self._sweep = target_sweep
                        self._use_kill_all = use_kill_all

                    @pyqtSlot()
                    def do_kill(self):
                        # For SweepQueue, use kill_all() to kill current AND clear queue
                        if self._use_kill_all and hasattr(self._sweep, "kill_all"):
                            self._sweep.kill_all()
                        else:
                            self._sweep.kill()

                proxy = _KillProxy(sweep, use_kill_all=is_sweep_queue)
                proxy.moveToThread(sweep.thread())
                sweep._mcp_kill_proxy = proxy

            # If sweep is already not running, try to call kill via Qt proxy to release resources
            # For SweepQueue, check effective_state; for BaseSweep, check state
            if is_sweep_queue:
                is_not_running = previous_state not in ("running", "pending", "paused")
            else:
                is_not_running = previous_state not in SWEEP_STATE_RUNNING

            if is_not_running:
                try:
                    # Queue the kill on the sweep's Qt thread
                    QMetaObject.invokeMethod(proxy, "do_kill", QueuedConnection)
                    # Brief wait for Qt to process
                    await asyncio.sleep(0.1)
                    return {
                        "success": True,
                        "sweep_name": var_name,
                        "sweep_type": type(sweep).__name__,
                        "previous_state": previous_state,
                        "new_state": previous_state,
                        "message": f"Sweep '{var_name}' killed to release resources (was in state: {previous_state}).",
                    }
                except Exception as e:
                    # Qt kill failed - check state and report
                    logger.debug(f"Qt kill on stopped sweep '{var_name}' failed: {e}")
                    progress_state = _get_progress_state(sweep)
                    if not progress_state:
                        # Can't access state, assume sweep is stuck
                        return {
                            "success": False,
                            "sweep_name": var_name,
                            "error": f"Cannot access progress state for '{var_name}' after kill attempt",
                        }

                    current_state = progress_state.state.value
                    if current_state not in SWEEP_STATE_RUNNING:
                        # Sweep is still stopped - no need to kill
                        return {
                            "success": True,
                            "sweep_name": var_name,
                            "sweep_type": type(sweep).__name__,
                            "previous_state": previous_state,
                            "new_state": current_state,
                            "message": f"Sweep '{var_name}' was already stopped (state: {current_state}). No kill needed.",
                        }
                    else:
                        # Sweep somehow started running - fall through to retry loop
                        logger.warning(
                            "Sweep '%s' state changed from '%s' to '%s' during kill attempt",
                            var_name,
                            previous_state,
                            current_state,
                        )
                        previous_state = current_state

            # Retry kill with verification - Qt event loop may need time to process
            max_retries = 5
            retry_delay = 0.3  # seconds between retries

            for attempt in range(max_retries):
                # Queue the kill on the sweep's Qt thread
                QMetaObject.invokeMethod(proxy, "do_kill", QueuedConnection)

                # Wait for the kill to take effect
                await asyncio.sleep(retry_delay)

                # Check if sweep actually stopped
                # For SweepQueue, also verify queue is empty (kill_all clears it)
                progress_state = _get_progress_state(sweep)
                if is_sweep_queue:
                    # Always use status() for SweepQueue to check queue_length
                    queue_status = _get_sweep_queue_status(sweep)
                    if queue_status:
                        new_state = queue_status.get("effective_state", "idle")
                        queue_length = queue_status.get("queue_length", 0)
                        # Stopped means: not active AND queue is empty
                        is_stopped = (
                            new_state not in ("running", "pending", "paused")
                            and queue_length == 0
                        )
                    elif progress_state:
                        # Fallback: use progress_state
                        new_state = progress_state.state.value
                        is_stopped = new_state not in ("running", "ramping")
                    else:
                        new_state = "idle"
                        is_stopped = True
                elif progress_state:
                    new_state = progress_state.state.value
                    is_stopped = new_state not in SWEEP_STATE_RUNNING
                else:
                    # BaseSweep without progress_state - shouldn't happen
                    new_state = "unknown"
                    is_stopped = False

                if is_stopped:
                    # Kill successful - sweep is no longer running
                    return {
                        "success": True,
                        "sweep_name": var_name,
                        "sweep_type": type(sweep).__name__,
                        "previous_state": previous_state,
                        "new_state": new_state,
                        "message": f"Sweep '{var_name}' killed successfully. Resources released.",
                        "warning": "UNSAFE: Sweep was terminated. Instruments may need re-initialization.",
                    }

                # Log retry attempt
                if attempt < max_retries - 1:
                    logger.warning(
                        "Kill attempt %d/%d for sweep '%s' - still in state '%s', retrying...",
                        attempt + 1,
                        max_retries,
                        var_name,
                        new_state,
                    )

            # All retries exhausted - sweep is still running
            progress_state = _get_progress_state(sweep)
            if progress_state:
                final_state = progress_state.state.value
            elif is_sweep_queue:
                queue_status = _get_sweep_queue_status(sweep)
                final_state = (
                    queue_status.get("effective_state", "unknown")
                    if queue_status
                    else "unknown"
                )
            else:
                final_state = "unknown"
            logger.error(
                "Failed to kill sweep '%s' after %d attempts. State: %s -> %s",
                var_name,
                max_retries,
                previous_state,
                final_state,
            )

            return {
                "success": False,
                "sweep_name": var_name,
                "sweep_type": type(sweep).__name__,
                "previous_state": previous_state,
                "new_state": final_state,
                "error": (
                    f"Sweep '{var_name}' is still running after {max_retries} kill attempts. "
                    f"State: {final_state}. The sweep may be stuck or unresponsive. "
                    "Try restarting the Jupyter kernel to force termination."
                ),
                "warning": "UNSAFE: Kill command was sent but sweep did not stop.",
            }

        except Exception as e:
            logger.error(f"Error killing sweep '{var_name}': {e}")
            return {
                "success": False,
                "sweep_name": var_name,
                "error": str(e),
            }

    async def kill_all_sweeps(self) -> Dict[str, Any]:
        """Kill all MeasureIt sweeps in the namespace to release resources.

        Iterates through all sweep objects in the notebook namespace and
        kills each one. This is useful for cleanup after errors or when
        multiple sweeps need to be terminated.

        Returns:
            Dict containing:
                success: bool - True if all sweeps were killed successfully
                killed_count: int - number of sweeps that were killed
                results: Dict mapping variable names to individual kill results
                errors: List of error messages (if any sweeps failed to kill)
        """
        try:
            if BaseSweep is None:
                return {
                    "success": False,
                    "killed_count": 0,
                    "results": {},
                    "errors": ["MeasureIt library not available"],
                }

            # Find all sweep objects in namespace (skip queued sweeps)
            sweep_vars = []
            for var_name, var_value in self.namespace.items():
                if var_name.startswith("_"):
                    continue
                if _is_sweep(var_value):
                    # Skip individual sweeps (BaseSweep) managed by a SweepQueue
                    # Don't skip SweepQueue objects themselves
                    is_base_sweep = BaseSweep is not None and isinstance(
                        var_value, BaseSweep
                    )
                    if is_base_sweep:
                        progress_state = _get_progress_state(var_value)
                        if progress_state and getattr(
                            progress_state, "is_queued", False
                        ):
                            continue
                    sweep_vars.append(var_name)

            if not sweep_vars:
                return {
                    "success": True,
                    "killed_count": 0,
                    "results": {},
                    "message": "No sweeps found in namespace",
                }

            # Kill each sweep
            results = {}
            errors = []
            killed_count = 0

            for var_name in sweep_vars:
                result = await self.kill_sweep(var_name)
                results[var_name] = result
                if result.get("success"):
                    killed_count += 1
                else:
                    errors.append(f"{var_name}: {result.get('error', 'Unknown error')}")

            return {
                "success": len(errors) == 0,
                "killed_count": killed_count,
                "total_sweeps": len(sweep_vars),
                "results": results,
                "errors": errors if errors else None,
                "message": f"Killed {killed_count}/{len(sweep_vars)} sweeps",
            }

        except Exception as e:
            logger.error(f"Error killing all sweeps: {e}")
            return {
                "success": False,
                "killed_count": 0,
                "results": {},
                "errors": [str(e)],
            }

    async def wait_for_sweep(
        self,
        var_name: str,
        timeout: Optional[float] = None,
        kill: bool = True,
        progress_callback: Optional[ProgressCallback] = None,
    ) -> Dict[str, Any]:
        """Wait for a measureit sweep with a given variable name to finish.

        Waits for the sweep with the given name to stop running and optionally
        kills it to release hardware resources.

        When kill=True (default), sweeps are killed in the following cases:
        - If already in error state when called: killed immediately
        - If already finished (done state) when called: killed immediately
        - If it enters error state while waiting: killed immediately
        - When it finishes successfully: killed to release hardware resources

        When kill=False, the sweep is not killed after waiting.

        Note: If timeout is reached, the sweep is NOT killed (still running);
        a kill_suggestion is provided for the user to decide.

        Args:
            var_name: Name of the sweep variable to wait for.
            timeout: Maximum time to wait in seconds. If None, wait indefinitely.
            kill: If True (default), kill the sweep after waiting to release resources.
                  If False, skip killing the sweep.

        Returns:
            Dict containing:
                sweep: Dict of information about the sweep as in get_measureit_status, or None if no
                running sweep with this name exists.
                error: str (if any error occurred, including timeout or sweep error)
                timed_out: bool (True if timeout was reached)
                sweep_error: bool (True if sweep ended with error state)
                killed: bool (True if sweep was killed successfully)
                kill_result: Dict with result of killing the sweep
        """
        status = await self.get_measureit_status()
        if status.get("error"):
            return {"sweep": None, "error": status["error"]}
        target = status["sweeps"].get(var_name)

        if not target:
            return {"sweep": None}

        # Check if sweep is already in error state
        # For SweepQueue, also check "stopped" state (indicates previous error)
        state = target["state"]
        is_sweep_queue = target.get("type") == "SweepQueue"
        if is_sweep_queue:
            is_error = state in ("error", "stopped")
        else:
            is_error = state == SWEEP_STATE_ERROR

        if is_error:
            error_msg = target.get(
                "error_message", f"Sweep '{var_name}' is already in error state"
            )
            result = {
                "sweep": target,
                "error": error_msg,
                "sweep_error": True,
            }
            if kill:
                logger.warning(
                    "Sweep '%s' already in error state: %s. Killing sweep.",
                    var_name,
                    error_msg,
                )
                kill_result = await self.kill_sweep(var_name)
                result["killed"] = kill_result.get("success", False)
                result["kill_result"] = kill_result
            else:
                result["killed"] = False
            return result

        # If SweepQueue is "pending" (has items but not started), return immediately
        # because there's no active sweep to wait for - the queue needs to be started first
        if is_sweep_queue and state == "pending":
            return {
                "sweep": target,
                "not_started": True,
                "killed": False,
                "message": (
                    f"SweepQueue '{var_name}' has {target.get('queue_length', 0)} items "
                    f"queued but has not been started. Call {var_name}.start() to start."
                ),
            }

        # If sweep is already done (not running, not error)
        if is_sweep_queue:
            # "paused" is considered done - return so caller can decide what to do
            is_done = state != "running"
        else:
            is_done = state not in SWEEP_STATE_RUNNING

        if is_done:
            result = {"sweep": target}
            if kill:
                logger.info(
                    "Sweep '%s' already done (state: %s). Killing to release resources.",
                    var_name,
                    state,
                )
                kill_result = await self.kill_sweep(var_name)
                result["killed"] = kill_result.get("success", False)
                result["kill_result"] = kill_result
            else:
                result["killed"] = False
            return result

        start_time = time.time()
        poll_count = 0
        while True:
            await asyncio.sleep(WAIT_DELAY)
            poll_count += 1

            # Send progress notification to keep client connection alive
            if progress_callback:
                elapsed = time.time() - start_time
                total = timeout if timeout else 0
                try:
                    await progress_callback(
                        elapsed, total, f"Waiting for '{var_name}'..."
                    )
                except Exception:
                    pass  # Don't let notification failures break the wait

            # Refresh status FIRST, then check timeout (so we have accurate state)
            status = await self.get_measureit_status()
            if status.get("error"):
                return {"sweep": target, "error": status["error"]}

            # Update target with latest state (None if sweep disappeared)
            target = status["sweeps"].get(var_name)

            if not target:
                # Sweep no longer in namespace
                return {"sweep": None}

            # Check for error state FIRST - kill before returning timeout
            # For SweepQueue, also check "stopped" state (indicates previous error)
            loop_state = target["state"]
            loop_is_sweep_queue = target.get("type") == "SweepQueue"
            if loop_is_sweep_queue:
                loop_is_error = loop_state in ("error", "stopped")
            else:
                loop_is_error = loop_state == SWEEP_STATE_ERROR

            if loop_is_error:
                error_msg = target.get("error_message", "Sweep ended with error")
                result = {
                    "sweep": target,
                    "error": error_msg,
                    "sweep_error": True,
                }
                if kill:
                    logger.warning(
                        "Sweep '%s' entered error state: %s. Killing sweep.",
                        var_name,
                        error_msg,
                    )
                    kill_result = await self.kill_sweep(var_name)
                    result["killed"] = kill_result.get("success", False)
                    result["kill_result"] = kill_result
                else:
                    result["killed"] = False
                return result

            # Check timeout AFTER error handling
            if timeout is not None and (time.time() - start_time) > timeout:
                return {
                    "sweep": target,
                    "error": f"Timeout after {timeout}s waiting for sweep '{var_name}' to complete",
                    "timed_out": True,
                    "kill_suggestion": f"Use measureit_kill_sweep('{var_name}') to stop the sweep and release resources",
                }

            # Check if sweep is done
            # For SweepQueue: only "running" is active
            # "pending" means queue has items but isn't started - return immediately
            # "idle", "paused", "error", "stopped" = done (return so caller can decide)
            state = target["state"]
            is_sweep_queue = target.get("type") == "SweepQueue"
            if is_sweep_queue:
                # Only "running" is considered active for SweepQueue
                is_active = state == "running"
            else:
                # Regular sweep just checks running states
                is_active = state in SWEEP_STATE_RUNNING
            if not is_active:
                break

        # Sweep finished successfully
        result = {"sweep": target}
        if kill:
            kill_result = await self.kill_sweep(var_name)
            result["killed"] = kill_result.get("success", False)
            result["kill_result"] = kill_result
        else:
            result["killed"] = False

        return result

    async def wait_for_all_sweeps(
        self,
        timeout: Optional[float] = None,
        kill: bool = True,
        progress_callback: Optional[ProgressCallback] = None,
    ) -> Dict[str, Any]:
        """Wait until all running measureit sweeps finish.

        Waits until all currently running sweeps have stopped running and optionally
        kills them to release hardware resources.

        When kill=True (default), sweeps are killed in the following cases:
        - If already in error state when called: killed immediately (but waiting continues)
        - If already finished (done state) when called: killed immediately
        - If they enter error state while waiting: killed immediately (but waiting continues)
        - When they finish successfully: killed to release hardware resources

        When kill=False, sweeps are not killed after waiting.

        Note: If timeout is reached, sweeps are NOT killed (still running);
        a kill_suggestion is provided for the user to decide.

        Args:
            timeout: Maximum time to wait in seconds. If None, wait indefinitely.
            kill: If True (default), kill sweeps after waiting to release resources.
                  If False, skip killing sweeps.

        Returns:
            Dict containing:
                sweeps: Dict mapping variable names to Dicts of information about the initially running sweeps as in get_measureit_status, empty if no sweeps were running.
                error: str (if any error occurred, including timeout or sweep error)
                timed_out: bool (True if timeout was reached)
                sweep_error: bool (True if any sweep ended with error state)
                errored_sweeps: list of sweep names that had errors (if any)
                killed: bool (True if sweeps were killed successfully)
                kill_results: Dict mapping sweep names to kill results
        """
        try:
            status = await self.get_measureit_status()
        except Exception as exc:  # pragma: no cover - defensive
            logger.error("wait_for_all_sweeps failed to get status: %s", exc)
            return {"sweeps": None, "error": str(exc)}

        if status.get("error"):
            return {"sweeps": None, "error": status["error"]}

        sweeps = status["sweeps"]

        # Track all sweeps we'll monitor and their final results
        all_kill_results = {}
        error_messages = []
        errored_sweeps = []

        # Helper to categorize sweep state
        def _categorize_sweep(sweep_info):
            """Returns 'running', 'pending', 'error', or 'done'."""
            state = sweep_info["state"]
            is_sweep_queue = sweep_info.get("type") == "SweepQueue"
            if is_sweep_queue:
                # SweepQueue: only "running" is active
                # "pending" means queue has items but isn't started - separate category
                # "paused" is "done" - return so caller can decide what to do
                if state == "running":
                    return "running"
                elif state == "pending":
                    return "pending"
                elif state in ("error", "stopped"):
                    return "error"
                else:  # "idle", "paused"
                    return "done"
            else:
                # Regular sweep
                if state in SWEEP_STATE_RUNNING:
                    return "running"
                elif state == SWEEP_STATE_ERROR:
                    return "error"
                else:
                    return "done"

        # Handle sweeps already in error state
        already_errored = {
            k: v for k, v in sweeps.items() if _categorize_sweep(v) == "error"
        }
        for name, sweep_info in already_errored.items():
            error_msg = sweep_info.get(
                "error_message", f"Sweep '{name}' is already in error state"
            )
            error_messages.append(f"{name}: {error_msg}")
            errored_sweeps.append(name)
            if kill:
                logger.warning(
                    "Sweep '%s' already in error state: %s. Killing sweep.",
                    name,
                    error_msg,
                )
                all_kill_results[name] = await self.kill_sweep(name)

        # Handle sweeps already in done state
        already_done = {
            k: v for k, v in sweeps.items() if _categorize_sweep(v) == "done"
        }
        if kill:
            for name in already_done.keys():
                logger.info(
                    "Sweep '%s' already done. Killing to release resources.", name
                )
                all_kill_results[name] = await self.kill_sweep(name)

        # Handle SweepQueues in pending state (items queued but not started)
        # Don't kill these - just report them with not_started info
        pending_queues = {
            k: v for k, v in sweeps.items() if _categorize_sweep(v) == "pending"
        }

        initial_running = {
            k: v for k, v in sweeps.items() if _categorize_sweep(v) == "running"
        }

        # If no running sweeps, return results from already-handled sweeps
        if not initial_running:
            if not already_errored and not already_done and not pending_queues:
                return {"sweeps": None}
            # Return combined results
            all_sweeps = {**already_errored, **already_done, **pending_queues}
            result = {"sweeps": all_sweeps}
            if kill and all_kill_results:
                all_killed = all(
                    r.get("success", False) for r in all_kill_results.values()
                )
                result["killed"] = all_killed
                result["kill_results"] = all_kill_results
            else:
                result["killed"] = False
            if pending_queues:
                result["not_started"] = list(pending_queues.keys())
                result["not_started_message"] = (
                    "Some SweepQueues have items queued but have not been started. "
                    "Call .start() on them to begin processing."
                )
            if errored_sweeps:
                result["error"] = "; ".join(error_messages)
                result["sweep_error"] = True
                result["errored_sweeps"] = errored_sweeps
            return result

        start_time = time.time()
        while True:
            await asyncio.sleep(WAIT_DELAY)

            # Send progress notification to keep client connection alive
            if progress_callback:
                elapsed = time.time() - start_time
                total = timeout if timeout else 0
                running_names = [
                    k
                    for k, v in initial_running.items()
                    if _categorize_sweep(v) == "running"
                ]
                try:
                    await progress_callback(
                        elapsed, total, f"Waiting for {len(running_names)} sweep(s)..."
                    )
                except Exception:
                    pass

            # Refresh status FIRST, then check timeout (so we have accurate state)
            status = await self.get_measureit_status()

            if status.get("error"):
                return {"sweeps": initial_running, "error": status["error"]}

            current_sweeps = status["sweeps"]
            still_running = False
            loop_errored = []

            # Update initial_running with latest state and detect errors
            for k in initial_running.keys():
                if k in current_sweeps:
                    initial_running[k] = current_sweeps[k]
                    sweep_info = current_sweeps[k]
                    state = sweep_info["state"]
                    is_sweep_queue = sweep_info.get("type") == "SweepQueue"

                    # Determine if sweep is active
                    if is_sweep_queue:
                        # SweepQueue: only "running" is active
                        # "pending" means queue has items but isn't started
                        is_active = state == "running"
                        is_error = state in ("error", "stopped")
                    else:
                        # Regular sweep
                        is_active = state in SWEEP_STATE_RUNNING
                        is_error = state == SWEEP_STATE_ERROR

                    if is_active:
                        still_running = True
                    elif is_error:
                        loop_errored.append(k)

            # Handle errored sweeps (do this BEFORE timeout check)
            if loop_errored:
                for name in loop_errored:
                    if name not in errored_sweeps:  # Don't double-process
                        sweep_info = initial_running.get(name, {})
                        msg = sweep_info.get(
                            "error_message", f"Sweep '{name}' ended with error"
                        )
                        error_messages.append(f"{name}: {msg}")
                        errored_sweeps.append(name)
                        if kill:
                            logger.warning(
                                "Sweep '%s' entered error state: %s. Killing sweep.",
                                name,
                                msg,
                            )
                            all_kill_results[name] = await self.kill_sweep(name)

            # Check timeout AFTER processing errors
            if timeout is not None and (time.time() - start_time) > timeout:
                running_names = [
                    k
                    for k, v in initial_running.items()
                    if _categorize_sweep(v) == "running"
                ]

                # Combine all sweeps and include already-killed results
                all_sweeps = {
                    **already_errored,
                    **already_done,
                    **pending_queues,
                    **initial_running,
                }
                result = {
                    "sweeps": all_sweeps,
                    "error": f"Timeout after {timeout}s waiting for sweeps to complete",
                    "timed_out": True,
                    "kill_suggestion": (
                        ", ".join(
                            [f"measureit_kill_sweep('{n}')" for n in running_names]
                        )
                        if running_names
                        else None
                    ),
                }
                if all_kill_results:
                    result["kill_results"] = all_kill_results
                    result["killed"] = all(
                        r.get("success", False) for r in all_kill_results.values()
                    )
                else:
                    result["killed"] = False
                if pending_queues:
                    result["not_started"] = list(pending_queues.keys())
                    result["not_started_message"] = (
                        "Some SweepQueues have items queued but have not been started. "
                        "Call .start() on them to begin processing."
                    )
                if errored_sweeps:
                    result["sweep_error"] = True
                    result["errored_sweeps"] = errored_sweeps
                    result["error"] = f"Timeout after {timeout}s; errors: " + "; ".join(
                        error_messages
                    )
                return result

            if not still_running:
                break

        # Kill all finished sweeps to release hardware resources
        if kill:
            for name in initial_running.keys():
                if name not in all_kill_results:  # Don't double-kill errored ones
                    all_kill_results[name] = await self.kill_sweep(name)

        # Combine all results
        all_sweeps = {
            **already_errored,
            **already_done,
            **pending_queues,
            **initial_running,
        }
        result = {"sweeps": all_sweeps}

        if kill and all_kill_results:
            all_killed = all(r.get("success", False) for r in all_kill_results.values())
            result["killed"] = all_killed
            result["kill_results"] = all_kill_results
        else:
            result["killed"] = False

        if pending_queues:
            result["not_started"] = list(pending_queues.keys())
            result["not_started_message"] = (
                "Some SweepQueues have items queued but have not been started. "
                "Call .start() on them to begin processing."
            )

        if errored_sweeps:
            result["error"] = "; ".join(error_messages)
            result["sweep_error"] = True
            result["errored_sweeps"] = errored_sweeps

        return result
