"""Module for testing the PandaBoxCSAXS and PandaBoxOMNY devices."""

# pylint: skip-file
from __future__ import annotations

import threading
from unittest import mock

import numpy as np
import pytest
from ophyd import Staged
from ophyd_devices.devices.panda_box.panda_box import PandaState
from ophyd_devices.interfaces.base_classes.psi_device_base import DeviceStoppedError
from pandablocks.responses import FrameData

from csaxs_bec.devices.panda_box.panda_box import PandaBoxCSAXS
from csaxs_bec.devices.panda_box.panda_box_omny import PandaBoxOMNY
from csaxs_bec.devices.utils.utils import fetch_scan_info


def _make_frame_data(n_rows: int, start: int = 0) -> FrameData:
    """Build a synthetic FrameData batch with `n_rows` rows for the frame_counter /
    gate_detector_active signals used by the raw-stream tests below."""
    data = np.array(
        [(np.float64(i), np.float64(i % 2)) for i in range(start, start + n_rows)],
        dtype=[("COUNTER1.OUT.Value", "<f8"), ("PCAP.BITS0.Value", "<f8")],
    )
    return FrameData(data)


@pytest.fixture
def panda_omny():
    dev_name = "panda_omny"
    dev = PandaBoxOMNY(
        name=dev_name,
        host="omny-panda-box.psi.ch",
        signal_alias={
            "FMC_IN.VAL1.Min": "cap_voltage_fzp_y_min",
            "FMC_IN.VAL1.Max": "cap_voltage_fzp_y_max",
            "FMC_IN.VAL1.Mean": "cap_voltage_fzp_y_mean",
            "FMC_IN.VAL2.Min": "cap_voltage_fzp_x_min",
            "FMC_IN.VAL2.Max": "cap_voltage_fzp_x_max",
            "FMC_IN.VAL2.Mean": "cap_voltage_fzp_x_mean",
        },
    )
    yield dev


@pytest.fixture
def panda_omny_raw_stream():
    dev_name = "panda_omny_raw"
    dev = PandaBoxOMNY(
        name=dev_name,
        host="omny-panda-box.psi.ch",
        signal_alias={
            "PCAP.BITS0.Value": "gate_detector_active",
            "COUNTER1.OUT.Value": "frame_counter",
        },
        raw_stream_mode=True,
        raw_stream_flush_row_count=5,
        # Effectively disable the time-based flush so tests are deterministic and only
        # exercise the row-count threshold.
        raw_stream_flush_interval=1000,
    )
    yield dev


@pytest.fixture
def panda_csaxs():
    dev_name = "panda_csaxs"
    dev = PandaBoxCSAXS(name=dev_name, host="csaxs-panda-box.psi.ch")
    yield dev


def test_panda_omny(panda_omny):
    assert panda_omny.name == "panda_omny"
    assert panda_omny.host == "omny-panda-box.psi.ch"

    all_signal_names = [name for name, _ in panda_omny.data.signals]
    # Check that the signal aliases are correctly set up
    assert "cap_voltage_fzp_y_min" in all_signal_names
    assert "cap_voltage_fzp_y_max" in all_signal_names
    assert "cap_voltage_fzp_y_mean" in all_signal_names
    assert "cap_voltage_fzp_x_min" in all_signal_names
    assert "cap_voltage_fzp_x_max" in all_signal_names
    assert "cap_voltage_fzp_x_mean" in all_signal_names

    # Check that the original signal names are not present
    assert "FMC_IN.VAL1.Min" not in all_signal_names
    assert "FMC_IN.VAL1.Max" not in all_signal_names
    assert "FMC_IN.VAL1.Mean" not in all_signal_names
    assert "FMC_IN.VAL2.Min" not in all_signal_names
    assert "FMC_IN.VAL2.Max" not in all_signal_names
    assert "FMC_IN.VAL2.Mean" not in all_signal_names

    assert panda_omny._acquisition_group == "burst"
    assert panda_omny._timeout_on_completed == 10


@pytest.mark.parametrize(
    "scan_type, frames_per_trigger, expected_acquisition_group",
    [
        ("hardware_triggered", 1, "fly"),
        ("hardware_triggered", 5, "fly"),
        ("software_triggered", 10, "burst"),
        ("software_triggered", 1, "monitored"),  # Default case
    ],
)
def test_panda_omny_stage(panda_omny, scan_type, frames_per_trigger, expected_acquisition_group):
    # Check that the stage signal is present and has the correct PV
    assert len(panda_omny._status_callbacks) == 0

    panda_omny.scan_info.msg.info["scan_type"] = scan_type
    panda_omny.scan_info.msg.info["frames_per_trigger"] = frames_per_trigger
    panda_omny.stage()

    assert panda_omny._acquisition_group == expected_acquisition_group
    assert panda_omny.staged == Staged.yes


def test_panda_omny_complete(panda_omny):
    """Test the on_complete method of the PandaBoxCSAXS device."""
    panda_omny.scan_info.msg.info["num_points"] = 1
    panda_omny.scan_info.msg.info["frames_per_trigger"] = 1
    panda_omny.scan_parameters = fetch_scan_info(panda_omny.scan_info)

    panda_omny._timeout_on_completed = 0.5  # Set a short timeout for testing

    def _mock_return_captured(*args, **kwargs):
        return ["=0"]

    # Timeout Error on complete
    with (
        mock.patch.object(panda_omny, "send_raw", side_effect=_mock_return_captured),
        mock.patch.object(panda_omny, "_disarm", return_value=None) as mock_disarm,
    ):
        status = panda_omny.on_complete()
        assert status.done is False
        assert status.success is False

        with pytest.raises(TimeoutError):
            status.wait(timeout=4)
        mock_disarm.assert_called_once()

    # Successful complete
    panda_omny._timeout_on_completed = 5
    with (
        mock.patch.object(panda_omny, "send_raw", side_effect=[["=0"], ["=0"], ["=1"]]),
        mock.patch.object(panda_omny, "_disarm", return_value=None) as mock_disarm,
    ):
        status = panda_omny.on_complete()
        assert status.done is False
        assert status.success is False

        status.wait(timeout=4)
        mock_disarm.assert_called_once()
        assert status.done is True
        assert status.success is True


def test_panda_omny_default_raw_stream_mode_is_off(panda_omny):
    """New devices must default to legacy behavior unless raw_stream_mode is explicitly set."""
    assert panda_omny.raw_stream_mode is False


def test_panda_omny_legacy_receive_frame_data_is_unbuffered(panda_omny):
    """Backward-compatibility regression: without raw_stream_mode, each FrameData must still
    be published immediately, exactly as before this feature was added."""
    with mock.patch.object(panda_omny.data, "put") as mock_put:
        for i in range(3):
            panda_omny._receive_frame_data(_make_frame_data(2, start=i * 2))
        assert mock_put.call_count == 3


def test_panda_omny_raw_stream_acquisition_group(panda_omny_raw_stream):
    """raw_stream_mode must select the free_running acquisition group regardless of
    scan_type/frames_per_trigger."""
    panda_omny_raw_stream.scan_info.msg.info["scan_type"] = "software_triggered"
    panda_omny_raw_stream.scan_info.msg.info["frames_per_trigger"] = 10
    panda_omny_raw_stream.stage()

    assert panda_omny_raw_stream._acquisition_group == "free_running"
    assert panda_omny_raw_stream.staged == Staged.yes


def test_panda_omny_raw_stream_coalesces_by_row_count(panda_omny_raw_stream):
    """Multiple small FrameData batches must be coalesced into a single .data.put() call
    once the row-count threshold is reached, not published one-by-one."""
    dev = panda_omny_raw_stream
    with mock.patch.object(dev.data, "put") as mock_put:
        dev._receive_frame_data(_make_frame_data(2, start=0))
        dev._receive_frame_data(_make_frame_data(2, start=2))
        mock_put.assert_not_called()  # 4 rows buffered so far, threshold is 5

        dev._receive_frame_data(_make_frame_data(2, start=4))
        mock_put.assert_called_once()  # 6 rows >= 5, flush triggered

        out = mock_put.call_args.args[0]
        assert [v for v in out["frame_counter"]["value"]] == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
        assert [v for v in out["gate_detector_active"]["value"]] == [0.0, 1.0, 0.0, 1.0, 0.0, 1.0]


def test_panda_omny_raw_stream_gate_bit_is_masked_from_shared_word(panda_omny_raw_stream):
    """PCAP.BITS0 is a 32-bit word shared with other captured bit_out signals; only bit 0
    (TTLIN1.VAL, confirmed via CAPTURE_WORD?/OFFSET? against the omny-panda hardware) is the
    detector-trigger gate. gate_detector_active must be masked down to that bit rather than
    passed through as the raw word."""
    dev = panda_omny_raw_stream
    data = np.array(
        # bit 0 set (gate active) plus other bits from something else sharing BITS0
        [(0.0, 5.0), (1.0, 4.0), (2.0, 3.0)],
        dtype=[("COUNTER1.OUT.Value", "<f8"), ("PCAP.BITS0.Value", "<f8")],
    )
    with mock.patch.object(dev.data, "put") as mock_put:
        dev._receive_frame_data(FrameData(data))
        dev._flush_raw_stream_buffer()

        out = mock_put.call_args.args[0]
        assert [v for v in out["gate_detector_active"]["value"]] == [1, 0, 1]


def test_panda_omny_raw_stream_flush_on_demand(panda_omny_raw_stream):
    """Buffered data below the threshold must still be flushed on demand (as on_complete/
    on_unstage do), and flushing an empty buffer must not publish again."""
    dev = panda_omny_raw_stream
    with mock.patch.object(dev.data, "put") as mock_put:
        dev._receive_frame_data(_make_frame_data(2, start=0))
        mock_put.assert_not_called()

        dev._flush_raw_stream_buffer()
        mock_put.assert_called_once()
        out = mock_put.call_args.args[0]
        assert len(out["frame_counter"]["value"]) == 2

        dev._flush_raw_stream_buffer()
        mock_put.assert_called_once()  # still 1: nothing left to flush


def test_panda_omny_unstage_flushes_pending_raw_stream_data(panda_omny_raw_stream):
    dev = panda_omny_raw_stream
    dev.stage()
    with (
        mock.patch.object(dev.data, "put") as mock_put,
        mock.patch.object(dev, "_disarm", return_value=None),
    ):
        dev._receive_frame_data(_make_frame_data(2, start=0))
        mock_put.assert_not_called()

        dev.unstage()
        mock_put.assert_called_once()


def test_panda_omny_complete_free_running_waits_for_frame_counter(panda_omny_raw_stream):
    """Completion in free-running mode waits for COUNTER1 (frame_counter) to reach the exact
    expected frame count -- PCAP.CAPTURED can't be used since it free-runs continuously off
    CLOCK1 and never naturally stops increasing while armed."""
    dev = panda_omny_raw_stream
    dev.scan_info.msg.info["num_points"] = 1
    dev.scan_info.msg.info["frames_per_trigger"] = 1
    dev.scan_parameters = fetch_scan_info(dev.scan_info)
    dev._timeout_on_completed = 5

    with (
        mock.patch.object(
            dev, "send_raw", side_effect=[["=0"], ["=0"], ["=1"]]
        ) as mock_send_raw,
        mock.patch.object(dev, "_disarm", return_value=None) as mock_disarm,
        mock.patch.object(dev, "_flush_raw_stream_buffer") as mock_flush,
    ):
        status = dev.on_complete()
        status.wait(timeout=4)
        assert status.done is True
        assert status.success is True
        mock_disarm.assert_called_once()
        mock_flush.assert_called_once()
        assert mock_send_raw.call_args.args[0] == "COUNTER1.OUT?"


def test_panda_omny_complete_free_running_timeout(panda_omny_raw_stream):
    """If the frame count never reaches the expected total, the overall timeout must still
    fire."""
    dev = panda_omny_raw_stream
    dev.scan_info.msg.info["num_points"] = 1
    dev.scan_info.msg.info["frames_per_trigger"] = 1
    dev.scan_parameters = fetch_scan_info(dev.scan_info)
    dev._timeout_on_completed = 0.3

    with (
        mock.patch.object(dev, "send_raw", side_effect=lambda *a, **k: ["=0"]),
        mock.patch.object(dev, "_disarm", return_value=None) as mock_disarm,
        mock.patch.object(dev, "_flush_raw_stream_buffer") as mock_flush,
    ):
        status = dev.on_complete()
        with pytest.raises(TimeoutError):
            status.wait(timeout=4)
        mock_disarm.assert_called_once()
        mock_flush.assert_called_once()


def _deliver_ready_after(dev, delay: float) -> threading.Timer:
    """Simulate the PandA data thread delivering the READY event `delay` seconds from now."""
    timer = threading.Timer(delay, dev._run_status_callbacks, args=(PandaState.READY,))
    timer.daemon = True
    timer.start()
    return timer


def test_panda_omny_pre_scan_arms_when_ready(panda_omny, capsys):
    dev = panda_omny
    dev.pre_scan_timeout = 2
    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda") as mock_reset,
    ):
        _deliver_ready_after(dev, 0.05)
        status = dev.on_pre_scan()
        status.wait(timeout=4)
        assert status.success is True
        mock_arm.assert_called_once()
        mock_reset.assert_not_called()
    assert "Forcing abort/reset" not in capsys.readouterr().out


def test_panda_omny_pre_scan_retries_once_then_succeeds(panda_omny, capsys):
    """No READY on the first attempt: the box is force-reset and arming retried. The retry is
    printed and logged."""
    dev = panda_omny
    dev.pre_scan_timeout = 0.2

    def _reset_then_ready():
        _deliver_ready_after(dev, 0.05)

    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda", side_effect=_reset_then_ready) as mock_reset,
        mock.patch("csaxs_bec.devices.panda_box.panda_box_omny.logger") as mock_logger,
    ):
        status = dev.on_pre_scan()
        status.wait(timeout=4)
        assert status.success is True
        mock_reset.assert_called_once()
        mock_arm.assert_called_once()
        assert any("Forcing abort/reset" in str(c) for c in mock_logger.warning.call_args_list)
        assert any("armed successfully on attempt 2" in str(c) for c in mock_logger.info.call_args_list)
    assert "Forcing abort/reset" in capsys.readouterr().out


def test_panda_omny_pre_scan_raises_after_retry_fails(panda_omny, capsys):
    dev = panda_omny
    dev.pre_scan_timeout = 0.1
    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda") as mock_reset,
        mock.patch("csaxs_bec.devices.panda_box.panda_box_omny.logger") as mock_logger,
    ):
        status = dev.on_pre_scan()
        with pytest.raises(RuntimeError, match="could not be armed after 2 attempts"):
            status.wait(timeout=4)
        mock_arm.assert_not_called()
        # once for the forced reset before the retry, once to leave the box clean at the end
        assert mock_reset.call_count == 2
        mock_logger.error.assert_called_once()
    assert capsys.readouterr().out.count("Forcing abort/reset") == 1
    # stale status callbacks of the failed attempts are removed again
    assert dev._status_callbacks == {}


def test_panda_omny_pre_scan_stop_is_not_retried(panda_omny, capsys):
    dev = panda_omny
    dev.pre_scan_timeout = 0.3
    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda") as mock_reset,
    ):
        status = dev.on_pre_scan()
        threading.Event().wait(0.05)
        dev._stop_stoppable_status_objects()
        with pytest.raises(DeviceStoppedError):
            status.wait(timeout=4)
        # give a wrongly retrying task thread the time to show up
        threading.Event().wait(0.5)
        mock_arm.assert_not_called()
        mock_reset.assert_not_called()
    assert "Forcing abort/reset" not in capsys.readouterr().out


def test_panda_omny_raw_stream_drops_stale_leading_frame_counter_rows(panda_omny_raw_stream):
    """PandA's COUNTER1.OUT is not cleared by Arm(), so the first raw samples of a new
    acquisition can briefly carry a stale frame_counter value left over from the previous
    acquisition's tail (confirmed 2026-09-21 against a real scan, S06313). Rows before the
    first observed frame_counter == 0 must be dropped, and every channel must be truncated
    at the same row so they stay aligned."""
    dev = panda_omny_raw_stream
    # Two stale rows carrying the previous acquisition's terminal counter value (167), then
    # the real acquisition starting at frame_counter 0.
    data = np.array(
        [(167.0, 1.0), (167.0, 0.0), (0.0, 1.0), (0.0, 0.0), (1.0, 1.0)],
        dtype=[("COUNTER1.OUT.Value", "<f8"), ("PCAP.BITS0.Value", "<f8")],
    )
    with mock.patch.object(dev.data, "put") as mock_put:
        dev._receive_frame_data(FrameData(data))
        dev._flush_raw_stream_buffer()

        out = mock_put.call_args.args[0]
        assert [v for v in out["frame_counter"]["value"]] == [0.0, 0.0, 1.0]
        assert [v for v in out["gate_detector_active"]["value"]] == [1, 0, 1]


def test_panda_omny_raw_stream_drops_whole_chunk_if_still_stale(panda_omny_raw_stream):
    """If a single FrameData batch is entirely stale rows (no frame_counter == 0 observed
    yet), it must be dropped in full rather than partially buffered, and the next batch
    must still be checked for the reset."""
    dev = panda_omny_raw_stream
    all_stale = np.array(
        [(167.0, 0.0), (167.0, 1.0)], dtype=[("COUNTER1.OUT.Value", "<f8"), ("PCAP.BITS0.Value", "<f8")]
    )
    real_start = np.array(
        [(0.0, 1.0), (1.0, 0.0)], dtype=[("COUNTER1.OUT.Value", "<f8"), ("PCAP.BITS0.Value", "<f8")]
    )
    with mock.patch.object(dev.data, "put") as mock_put:
        dev._receive_frame_data(FrameData(all_stale))
        assert dev._raw_stream_buffer == {}
        assert dev._raw_stream_seen_reset is False

        dev._receive_frame_data(FrameData(real_start))
        dev._flush_raw_stream_buffer()

        out = mock_put.call_args.args[0]
        assert [v for v in out["frame_counter"]["value"]] == [0.0, 1.0]


def test_panda_omny_raw_stream_passes_through_once_reset_seen(panda_omny_raw_stream):
    """After the first frame_counter == 0 has been observed, later rows must never be
    filtered again, even if a later point's real frame_counter happens to collide with the
    earlier stale value (e.g. point 167 in a scan with more than 167 points)."""
    dev = panda_omny_raw_stream
    with mock.patch.object(dev.data, "put"):
        dev._receive_frame_data(_make_frame_data(3, start=0))  # observes the reset at row 0
        assert dev._raw_stream_seen_reset is True
        dev._flush_raw_stream_buffer()

    # A later chunk whose frame_counter happens to read 167 again must pass through
    # untouched -- filtering only ever applies before the first reset.
    collide = np.array(
        [(167.0, 1.0), (168.0, 0.0)], dtype=[("COUNTER1.OUT.Value", "<f8"), ("PCAP.BITS0.Value", "<f8")]
    )
    with mock.patch.object(dev.data, "put") as mock_put:
        dev._receive_frame_data(FrameData(collide))
        dev._flush_raw_stream_buffer()
        out = mock_put.call_args.args[0]
        assert [v for v in out["frame_counter"]["value"]] == [167.0, 168.0]


def test_panda_omny_raw_stream_state_reset_on_stage(panda_omny_raw_stream):
    """on_stage must reopen the stale-value window (new acquisition) even if a previous
    acquisition already observed its reset."""
    dev = panda_omny_raw_stream
    dev._raw_stream_seen_reset = True
    dev._raw_stream_buffer = {"leftover": {"value": [1, 2, 3], "timestamp": 0.0}}

    dev.stage()

    assert dev._raw_stream_seen_reset is False
    assert dev._raw_stream_buffer == {}


def test_panda_omny_pre_scan_arm_resets_raw_stream_state(panda_omny):
    """Each Arm() resets COUNTER1's counting on the PandA side, so every successful arm in
    on_pre_scan -- not just the one implied by on_stage -- must reopen the stale-frame_counter
    window right before Arm() is sent. This matters most for a retried arm after a forced
    abort/reset, since on_stage's reset has already happened by then."""
    dev = panda_omny
    dev.pre_scan_timeout = 2
    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda"),
        mock.patch.object(dev, "_reset_raw_stream_state") as mock_reset_raw,
    ):
        _deliver_ready_after(dev, 0.05)
        status = dev.on_pre_scan()
        status.wait(timeout=4)
        assert status.success is True
        mock_arm.assert_called_once()
        mock_reset_raw.assert_called_once()


def test_panda_omny_arm_scan_allowlist_defaults_to_none(panda_omny):
    """A device constructed without arm_scan_allowlist must arm for every scan (today's
    behavior), both via the raw config attribute and the computed should-arm flag."""
    dev = panda_omny
    assert dev.arm_scan_allowlist is None
    assert dev._should_arm_panda is True


def test_panda_omny_stage_should_arm_stays_true_when_allowlist_none(panda_omny):
    """Backward-compatibility regression: staging with the default allowlist (None) must keep
    _should_arm_panda True regardless of the scan_name in scan_info."""
    dev = panda_omny
    dev.stage()
    assert dev._should_arm_panda is True


def test_panda_omny_stage_should_arm_true_when_scan_name_in_allowlist(panda_omny):
    dev = panda_omny
    dev.arm_scan_allowlist = ["flomni_fermat_scan"]
    dev.scan_info.msg.info["scan_name"] = "flomni_fermat_scan"
    dev.stage()
    assert dev._should_arm_panda is True


def test_panda_omny_stage_should_arm_false_when_scan_name_not_in_allowlist(panda_omny):
    dev = panda_omny
    dev.arm_scan_allowlist = ["flomni_fermat_scan"]
    dev.scan_info.msg.info["scan_name"] = "grid_scan"
    dev.stage()
    assert dev._should_arm_panda is False


def test_panda_omny_should_arm_recomputed_every_stage(panda_omny):
    """The same device instance is staged for many different scan types across a session; the
    should-arm decision must be recomputed fresh on every stage, not stick from a previous
    scan."""
    dev = panda_omny
    dev.arm_scan_allowlist = ["flomni_fermat_scan"]

    with mock.patch.object(dev, "_disarm"):
        dev.scan_info.msg.info["scan_name"] = "grid_scan"
        dev.stage()
        assert dev._should_arm_panda is False
        dev.unstage()

        dev.scan_info.msg.info["scan_name"] = "flomni_fermat_scan"
        dev.stage()
        assert dev._should_arm_panda is True


def test_panda_omny_pre_scan_skips_arm_when_should_arm_false(panda_omny):
    """When arm_scan_allowlist excludes the current scan, on_pre_scan must resolve immediately
    without sending Arm() or spinning up the retry task at all."""
    dev = panda_omny
    dev._should_arm_panda = False
    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_arm_with_retry") as mock_arm_with_retry,
        mock.patch.object(dev.task_handler, "submit_task") as mock_submit_task,
    ):
        status = dev.on_pre_scan()
        assert status.done is True
        assert status.success is True
        mock_arm.assert_not_called()
        mock_arm_with_retry.assert_not_called()
        mock_submit_task.assert_not_called()


def test_panda_omny_pre_scan_arms_when_should_arm_true_via_allowlist(panda_omny):
    """End-to-end: a scan_name matching arm_scan_allowlist still goes through the full arm/
    retry machinery exactly as with the default (None) allowlist."""
    dev = panda_omny
    dev.pre_scan_timeout = 2
    dev.arm_scan_allowlist = ["flomni_fermat_scan"]
    dev.scan_info.msg.info["scan_name"] = "flomni_fermat_scan"
    dev.stage()
    assert dev._should_arm_panda is True
    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda") as mock_reset,
    ):
        _deliver_ready_after(dev, 0.05)
        status = dev.on_pre_scan()
        status.wait(timeout=4)
        assert status.success is True
        mock_arm.assert_called_once()
        mock_reset.assert_not_called()


def test_panda_omny_complete_skips_wait_loop_when_should_arm_false(panda_omny):
    """When this device was never armed for the current scan, on_complete must resolve
    immediately without polling PandA or disarming it."""
    dev = panda_omny
    dev._should_arm_panda = False
    with (
        mock.patch.object(dev, "send_raw") as mock_send_raw,
        mock.patch.object(dev, "_disarm") as mock_disarm,
        mock.patch.object(dev.task_handler, "submit_task") as mock_submit_task,
    ):
        status = dev.on_complete()
        assert status.done is True
        assert status.success is True
        mock_send_raw.assert_not_called()
        mock_disarm.assert_not_called()
        mock_submit_task.assert_not_called()


def test_panda_omny_complete_skips_wait_loop_when_should_arm_false_raw_stream_mode(
    panda_omny_raw_stream,
):
    """Same as test_panda_omny_complete_skips_wait_loop_when_should_arm_false, but on the
    raw_stream_mode=True path that would otherwise dispatch to _on_complete_free_running."""
    dev = panda_omny_raw_stream
    dev._should_arm_panda = False
    with (
        mock.patch.object(dev, "send_raw") as mock_send_raw,
        mock.patch.object(dev, "_disarm") as mock_disarm,
        mock.patch.object(dev, "_flush_raw_stream_buffer") as mock_flush,
        mock.patch.object(dev.task_handler, "submit_task") as mock_submit_task,
    ):
        status = dev.on_complete()
        assert status.done is True
        assert status.success is True
        mock_send_raw.assert_not_called()
        mock_disarm.assert_not_called()
        mock_flush.assert_not_called()
        mock_submit_task.assert_not_called()


def test_panda_omny_raw_stream_signal_alias(panda_omny_raw_stream):
    all_signal_names = [name for name, _ in panda_omny_raw_stream.data.signals]
    assert "gate_detector_active" in all_signal_names
    assert "frame_counter" in all_signal_names
    assert "PCAP.BITS0.Value" not in all_signal_names
    assert "COUNTER1.OUT.Value" not in all_signal_names


def test_panda_csaxs(panda_csaxs):
    assert panda_csaxs.name == "panda_csaxs"
    assert panda_csaxs.host == "csaxs-panda-box.psi.ch"

    assert panda_csaxs._acquisition_group == "burst"
    assert panda_csaxs._timeout_on_completed == 10


@pytest.mark.parametrize(
    "scan_type, frames_per_trigger, expected_acquisition_group",
    [
        ("hardware_triggered", 1, "fly"),
        ("hardware_triggered", 5, "fly"),
        ("software_triggered", 10, "burst"),
        ("software_triggered", 1, "monitored"),  # Default case
    ],
)
def test_panda_csaxs_stage(panda_csaxs, scan_type, frames_per_trigger, expected_acquisition_group):
    """Test the on_stage method of the PandaBoxCSAXS device for different scan types and frames per trigger."""
    assert len(panda_csaxs._status_callbacks) == 0

    panda_csaxs.scan_info.msg.info["scan_type"] = scan_type
    panda_csaxs.scan_info.msg.info["frames_per_trigger"] = frames_per_trigger
    panda_csaxs.stage()

    assert panda_csaxs._acquisition_group == expected_acquisition_group
    assert panda_csaxs.staged == Staged.yes


def test_panda_csaxs_complete(panda_csaxs):
    """Test the on_complete method of the PandaBoxCSAXS device."""
    panda_csaxs.scan_info.msg.info["num_points"] = 1
    panda_csaxs.scan_info.msg.info["frames_per_trigger"] = 1
    panda_csaxs.scan_parameters = fetch_scan_info(panda_csaxs.scan_info)

    panda_csaxs._timeout_on_completed = 0.5  # Set a short timeout for testing

    def _mock_return_captured(*args, **kwargs):
        return ["=0"]

    # Timeout Error on complete
    with (
        mock.patch.object(panda_csaxs, "send_raw", side_effect=_mock_return_captured),
        mock.patch.object(panda_csaxs, "_disarm", return_value=None) as mock_disarm,
    ):
        status = panda_csaxs.on_complete()
        assert status.done is False
        assert status.success is False

        with pytest.raises(TimeoutError):
            status.wait(timeout=4)
        mock_disarm.assert_called_once()

    # Successful complete
    panda_csaxs._timeout_on_completed = 5
    with (
        mock.patch.object(panda_csaxs, "send_raw", side_effect=[["=0"], ["=0"], ["=1"]]),
        mock.patch.object(panda_csaxs, "_disarm", return_value=None) as mock_disarm,
    ):
        status = panda_csaxs.on_complete()
        assert status.done is False
        assert status.success is False

        status.wait(timeout=4)
        mock_disarm.assert_called_once()
        assert status.done is True
        assert status.success is True