"""River 2 protocol tests. Payload fixtures are synthetic, not hardware captures.

Offsets follow ha-ef-ble's DirectPdHeartbeatPack, DirectEmsDeltaHeartbeatPack,
and DirectBmsMDeltaHeartbeatPack. Fixtures do not use bmspy's field descriptors
so an accidental layout change cannot silently change the test data too.
"""

import asyncio
import importlib.util
import struct
import sys
from argparse import Namespace
from dataclasses import asdict
from itertools import permutations
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import pytest

from bmspy.classes import UPS
from bmspy.ecoflow_bt import (
    EcoFlowBTDevice,
    EcoFlowBTManager,
    EcoFlowPacket,
    _EncFrameHelper,
    _RawHeaderHelper,
    _State,
    _Type1Encryption,
    _Type7Encryption,
    _crc8,
    _crc16,
    _enc_packet_to_bytes,
    _parse_enc_frame,
    collect_data,
    initialise_device,
)
from bmspy.ecoflow_proto import decode_payload
from bmspy.ecoflow_river2 import decode_river2_payload

# Synthetic device identifiers, not values from a real device.
ADDRESS = "02:00:00:00:00:01"
SERIAL = "R601TEST00000001"


def device(serial=SERIAL, **kwargs):
    return EcoFlowBTDevice(ADDRESS, serial, "test-user", **kwargs)


def pd_payload(soc=73, watts_out=42, watts_in=160, remaining=87):
    data = bytearray(23)
    struct.pack_into("<BHHi", data, 14, soc, watts_out, watts_in, remaining)
    return bytes(data)


def ems_payload(charging=87, discharging=240, soc=None):
    data = bytearray(25 if soc is None else 30)
    data[0] = 1
    struct.pack_into("<II", data, 17, charging, discharging)
    if soc is not None:
        struct.pack_into("<f", data, 26, soc)
    return bytes(data)


def bms_payload(soc=73.25, temperature=29, watts_in=110, watts_out=0):
    data = bytearray(69)
    data[11] = 73
    data[20] = temperature
    struct.pack_into("<fIII", data, 53, soc, watts_in, watts_out, 87)
    return bytes(data)


def heartbeat(src, payload, cmd_id=0x02):
    return EcoFlowPacket(src, 0x21, 0x20, cmd_id, payload, version=2)


@pytest.mark.parametrize("serial,version", [
    (SERIAL, 2), ("R603XXXXXXXXXXXX", 2), ("R651TEST00000001", 3),
    ("unknown", 3),
])
def test_model_selects_packet_version_not_encryption(serial, version):
    dev = device(serial)
    assert dev._pkt_version == version
    assert dev._enc_type == 7


def test_explicit_packet_version_and_manager_defaults():
    assert device(packet_version=3)._pkt_version == 3
    manager = EcoFlowBTManager("test-user")
    assert manager.add_device(ADDRESS, SERIAL)._pkt_version == 2
    cfg = initialise_device(ADDRESS, SERIAL, "test-user")
    assert EcoFlowBTDevice(cfg.address, cfg.serial_number, cfg.user_id,
                          packet_version=cfg.packet_version)._pkt_version == 2


@pytest.mark.parametrize("version,overhead", [(2, 18), (3, 20)])
def test_packet_header_layout(version, overhead):
    pkt = EcoFlowPacket(0x21, 0x35, 0x35, 0x86, b"test", version=version)
    raw = pkt.to_bytes()
    assert len(raw) == overhead + 4
    assert raw[:4] == bytes([0xAA, version, 4, 0])
    assert raw[12:14] == b"\x21\x35"
    assert raw[overhead - 4:overhead - 2] == b"\x35\x86"
    if version == 3:
        assert raw[14:16] == b"\x01\x01"
    decoded = EcoFlowPacket.from_bytes(raw)
    assert decoded is not None
    assert decoded.payload == b"test"
    assert decoded.version == version
    assert decoded.cmd_set == 0x35
    assert decoded.cmd_id == 0x86


@pytest.mark.parametrize("raw", [b"", b"\xaa", b"\xaa\x02", b"\xaa\x02\x00\x00"])
def test_short_packet_is_rejected(raw):
    assert EcoFlowPacket.from_bytes(raw) is None


def test_bad_crc_and_inconsistent_length_are_rejected():
    raw = bytearray(heartbeat(2, pd_payload()).to_bytes())
    raw[-1] ^= 1
    assert EcoFlowPacket.from_bytes(raw) is None
    raw[2] += 1
    raw[4] = _crc8(raw[:4])
    raw[-2:] = struct.pack("<H", _crc16(raw[:-2]))
    assert EcoFlowPacket.from_bytes(raw) is None


def test_outer_frames_fragmentation_coalescing_and_embedded_prefix():
    # A marker inside an incomplete frame must not replace the real header.
    payload = b"abcd\x5a\x5a" + bytes(range(50))
    wire = _enc_packet_to_bytes(1, payload)
    second = _enc_packet_to_bytes(1, b"next")
    for split in range(len(wire) + 1):
        buf = bytearray()
        first, buf = _parse_enc_frame(wire[:split], buf)
        rest, buf = _parse_enc_frame(wire[split:] + second, buf)
        assert first + rest == [payload, b"next"]
        assert not buf


def test_outer_frame_discards_noise_and_bad_lengths():
    buf = bytearray()
    assert _parse_enc_frame(b"noise" * 100, buf)[0] == []
    assert not buf
    wire = _enc_packet_to_bytes(1, b"ok")
    for bad_length in [0, 1, 10001]:
        bad = b"\x5a\x5a\x10\x01" + struct.pack("<H", bad_length) + b"00"
        assert _parse_enc_frame(bad + wire, bytearray())[0] == [b"ok"]
    packets, buf = _parse_enc_frame(b"noise" + wire[:1], bytearray())
    assert packets == []
    assert _parse_enc_frame(wire[1:], buf)[0] == [b"ok"]
    corrupt = wire[:-1] + bytes([wire[-1] ^ 1])
    assert _parse_enc_frame(corrupt + wire, bytearray())[0] == [b"ok"]


@pytest.fixture
def enc_helper():
    pytest.importorskip("Crypto")
    return _EncFrameHelper(_Type7Encryption(bytes(range(16)), bytes(range(16, 32))))


@pytest.mark.parametrize("version", [2, 3])
def test_encrypted_frames_round_trip(enc_helper, version):
    pkt = EcoFlowPacket(2, 0x21, 0x20, 2, pd_payload(), version=version)
    wire = enc_helper.encode(pkt)
    for split in range(1, len(wire)):
        assert enc_helper.decode(wire[:split]) == []
        decoded = enc_helper.decode(wire[split:])
        assert len(decoded) == 1
        assert decoded[0].payload == pkt.payload
        assert decoded[0].version == version


def test_type1_version2_round_trip():
    pytest.importorskip("Crypto")
    helper = _RawHeaderHelper(_Type1Encryption(bytes(16), bytes(16)))
    pkt = heartbeat(2, pd_payload())
    wire = helper.encode(pkt)
    assert helper.decode(wire[:7]) == []
    assert helper.decode(wire[7:])[0].payload == pkt.payload


def test_numeric_decoder_uses_source_and_handles_short_firmware_payloads():
    pd = decode_river2_payload(2, 0x20, 2, pd_payload())
    ems = decode_river2_payload(3, 0x20, 2, ems_payload())
    assert pd["watts_in_sum"] == 160
    assert pd["watts_out_sum"] == 42
    assert pd["remain_time"] == 87
    assert ems["chg_remain_time"] == 87
    assert ems["dsg_remain_time"] == 240
    assert decode_river2_payload(2, 0x20, 2, pd_payload() + b"\x00") == {
        **pd, "beep_mode": 0,
    }
    assert decode_river2_payload(2, 0xFE, 0x15, pd_payload()) is None
    assert decode_river2_payload(9, 0x20, 2, pd_payload()) is None


@pytest.mark.parametrize("src,cmd_id,payload", [
    (2, 2, pd_payload()), (3, 2, ems_payload()), (3, 0x32, bms_payload()[:21]),
])
def test_truncated_core_heartbeat_is_rejected(src, cmd_id, payload):
    for length in range(len(payload)):
        assert decode_river2_payload(src, 0x20, cmd_id, payload[:length]) is None


def test_snapshot_contains_power_temperature_and_times():
    dev = device()
    updates = []
    dev.on_update(lambda: updates.append(dev.snapshot.ef_input_power_watts.raw_value))
    dev._record(heartbeat(2, pd_payload()))
    dev._record(heartbeat(3, ems_payload()))
    dev._record(heartbeat(3, bms_payload(), 0x32))
    snap = dev.snapshot
    assert updates == [160, 160, 160]
    assert snap.ef_input_power_watts.raw_value == 160
    assert snap.ef_input_power_watts.units == "W"
    assert snap.ef_output_power_watts.raw_value == 42
    assert snap.bms_capacity_charge_ratio.raw_value == 0.7325
    assert snap.bms_temperature_celsius.raw_values == {1: 29}
    assert snap.bms_temperature_celsius.units == "°C"
    assert snap.ef_remaining_time_minutes.raw_value == 87
    assert snap.ef_charge_remaining_time_minutes.raw_value == 87
    assert snap.ef_discharge_remaining_time_minutes.raw_value == 240
    assert snap.ef_charge_remaining_time_minutes.units == "min"
    assert snap.ef_charge_state.raw_value == 1
    assert snap.bms_charge_is_charging.raw_value is True
    assert snap.bms_charge_is_discharging.raw_value is False
    assert snap.ef_packets_received.raw_value == 3
    assert set(snap._packets) == {
        "ef_pkt_02_20_02", "ef_pkt_03_20_02", "ef_pkt_03_20_32",
    }
    # All metrics and source-specific packets survive the server/client format.
    serialized = {name: asdict(value) for name, value in snap.items()}
    assert dict(UPS.from_dict(serialized).items()) == dict(snap.items())


def test_zero_power_bypass_and_discharge():
    dev = device()
    dev._record(heartbeat(2, pd_payload(watts_in=42, watts_out=42)))
    assert dev.snapshot.bms_charge_is_charging is None
    dev._record(heartbeat(3, bms_payload(watts_in=0), 0x32))
    assert dev.snapshot.bms_charge_is_charging.raw_value is False
    dev._record(heartbeat(2, pd_payload(watts_in=0, watts_out=42)))
    dev._record(heartbeat(3, bms_payload(watts_in=0, watts_out=50), 0x32))
    assert dev.snapshot.ef_input_power_watts.raw_value == 0
    assert dev.snapshot.bms_charge_is_discharging.raw_value is True


def test_unavailable_times_clear_previous_values_and_zero_is_valid():
    dev = device()
    dev._record(heartbeat(2, pd_payload()))
    dev._record(heartbeat(3, ems_payload()))
    dev._record(heartbeat(2, pd_payload(remaining=-1)))
    dev._record(heartbeat(3, ems_payload(charging=0xFFFFFFFF, discharging=0)))
    assert dev.snapshot.ef_remaining_time_minutes is None
    assert dev.snapshot.ef_charge_remaining_time_minutes is None
    assert dev.snapshot.ef_discharge_remaining_time_minutes.raw_value == 0


def test_short_bms_has_temperature_without_inventing_charge_status():
    dev = device()
    dev._record(heartbeat(3, bms_payload()[:21], 0x32))
    assert dev.snapshot.bms_temperature_celsius.raw_values == {1: 29}
    assert dev.snapshot.bms_capacity_charge_ratio.raw_value == 0.73
    assert dev.snapshot.bms_charge_is_charging is None


@pytest.mark.parametrize("soc", [float("nan"), float("inf"), -1, 101])
def test_invalid_soc_is_not_exported(soc):
    dev = device()
    dev._record(heartbeat(3, bms_payload(soc=soc), 0x32))
    assert dev.snapshot.bms_capacity_charge_ratio is None


@pytest.mark.parametrize("order", list(permutations(["pd", "ems", "bms"])))
def test_precise_soc_is_independent_of_packet_order(order):
    dev = device()
    packets = {
        "pd": heartbeat(2, pd_payload(soc=89)),
        "ems": heartbeat(3, ems_payload(soc=88.91)),
        "bms": heartbeat(3, bms_payload(soc=88.918), 0x32),
    }
    for source in order:
        dev._record(packets[source])
    assert dev.snapshot.bms_capacity_charge_ratio.raw_value == pytest.approx(0.88918)
    # Source preference must not freeze the first precise reading.
    dev._record(heartbeat(3, bms_payload(soc=89.25), 0x32))
    dev._record(heartbeat(2, pd_payload(soc=89)))
    assert dev.snapshot.bms_capacity_charge_ratio.raw_value == 0.8925


def test_ems_float_is_a_fallback_when_bms_float_is_missing():
    dev = device()
    dev._record(heartbeat(3, ems_payload(soc=88.5)))
    dev._record(heartbeat(3, bms_payload()[:21], 0x32))
    dev._record(heartbeat(2, pd_payload(soc=89)))
    assert dev.snapshot.bms_capacity_charge_ratio.raw_value == 0.885
    dev._record(heartbeat(3, ems_payload(soc=89.5)))
    assert dev.snapshot.bms_capacity_charge_ratio.raw_value == 0.895


@pytest.mark.parametrize("soc", [0.0, 100.0])
def test_precise_soc_accepts_zero_and_full_charge(soc):
    dev = device()
    dev._record(heartbeat(3, bms_payload(soc=soc), 0x32))
    dev._record(heartbeat(2, pd_payload(soc=73)))
    assert dev.snapshot.bms_capacity_charge_ratio.raw_value == soc / 100


def test_invalid_float_does_not_block_valid_soc_fallback():
    dev = device()
    dev._record(heartbeat(3, bms_payload(soc=float("nan")), 0x32))
    dev._record(heartbeat(3, ems_payload(soc=float("inf"))))
    dev._record(heartbeat(2, pd_payload(soc=89)))
    assert dev.snapshot.bms_capacity_charge_ratio.raw_value == 0.89
    dev._record(heartbeat(3, bms_payload(soc=88.5), 0x32))
    dev._record(heartbeat(3, bms_payload(soc=float("nan")), 0x32))
    assert dev.snapshot.bms_capacity_charge_ratio.raw_value == 0.885


def test_soc_source_preference_resets_on_new_connection(monkeypatch):
    client = Mock(
        connect=AsyncMock(), disconnect=AsyncMock(), is_connected=True,
        start_notify=AsyncMock(), stop_notify=AsyncMock(),
    )
    monkeypatch.setitem(sys.modules, "bleak", SimpleNamespace(
        BleakClient=Mock(return_value=client),
    ))

    async def run():
        dev = device()
        dev._record(heartbeat(3, bms_payload(soc=88.5), 0x32))
        dev._auth_type7_init = AsyncMock()
        try:
            assert await dev.connect()
            dev._record(heartbeat(2, pd_payload(soc=90)))
            assert dev.snapshot.bms_capacity_charge_ratio.raw_value == 0.9
        finally:
            await dev.disconnect()
    asyncio.run(run())


@pytest.mark.parametrize("current", [-5, 0, 12])
def test_battery_pack_voltage_and_signed_current_units(current):
    bms = bytearray(bms_payload())
    ems = bytearray(ems_payload())
    struct.pack_into("<II", bms, 12, 13, current & 0xFFFFFFFF)
    struct.pack_into("<II", ems, 3, 13, current & 0xFFFFFFFF)
    struct.pack_into("<H", bms, 39, 3460)  # Cell voltage remains mV.
    struct.pack_into("<I", bms, 49, 20000)  # Charge-current limit remains mA.
    for cmd_id, payload, voltage_name, current_name in (
        (0x32, bytes(bms), "vol", "amp"),
        (0x02, bytes(ems), "chg_vol", "chg_amp"),
    ):
        values = decode_river2_payload(3, 0x20, cmd_id, payload)
        assert values[voltage_name] == 13
        assert values[current_name] == current
        lines = decode_payload(3, 0x20, cmd_id, payload)
        assert f"  {voltage_name}: 13 V" in lines
        assert f"  {current_name}: {current} A" in lines
    lines = decode_payload(3, 0x20, 0x32, bytes(bms))
    assert "  max_cell_vol: 3460 mV" in lines
    assert "  tag_chg_amp: 20000 mA" in lines


def test_mppt_voltage_setting_is_volts_not_millivolts():
    data = bytearray(75)
    struct.pack_into("<I", data, 8, 1122)
    struct.pack_into("<I", data, 61, 8000)
    struct.pack_into("<IB", data, 68, 230, 50)
    lines = decode_payload(5, 0x20, 2, bytes(data))
    assert "  cfg_ac_out_voltage: 230 V" in lines
    assert "  cfg_ac_out_freq: 50 Hz" in lines
    assert "  in_vol: 1122 mV" in lines
    assert "  cfg_dc_chg_current: 8000 mA" in lines


def test_unverified_inv_configuration_is_not_labeled_as_measured_frequency():
    data = bytearray(52)
    data[22] = 50
    data[51] = 2
    lines = decode_payload(4, 0x20, 2, bytes(data))
    assert "  inv_out_freq: 50 Hz" in lines
    assert "  cfg_ac_out_freq: 2" in lines


def test_capped_times_are_omitted_but_retained_in_diagnostics():
    dev = device()
    dev._record(heartbeat(2, pd_payload()))
    dev._record(heartbeat(3, ems_payload()))
    dev._record(heartbeat(2, pd_payload(remaining=5999)))
    dev._record(heartbeat(3, ems_payload(charging=5999, discharging=5999)))
    names = dict(dev.snapshot.items())
    assert "ef_remaining_time_minutes" not in names
    assert "ef_charge_remaining_time_minutes" not in names
    assert "ef_discharge_remaining_time_minutes" not in names
    raw = bytes.fromhex(names["ef_pkt_03_20_02"].info)
    assert "  chg_remain_time: 5999 min" in decode_payload(3, 0x20, 2, raw)


def test_switching_direction_clears_inactive_time_estimate():
    dev = device()
    dev._record(heartbeat(3, ems_payload(charging=69, discharging=5999)))
    assert dev.snapshot.ef_charge_remaining_time_minutes.raw_value == 69
    assert dev.snapshot.ef_discharge_remaining_time_minutes is None
    dev._record(heartbeat(3, ems_payload(charging=5999, discharging=184)))
    assert dev.snapshot.ef_charge_remaining_time_minutes is None
    assert dev.snapshot.ef_discharge_remaining_time_minutes.raw_value == 184


@pytest.mark.parametrize("minutes", [0, 10, 192, 5998, 6000])
def test_non_sentinel_times_are_not_adjusted(minutes):
    dev = device()
    dev._record(heartbeat(2, pd_payload(remaining=minutes)))
    dev._record(heartbeat(3, ems_payload(charging=minutes, discharging=minutes)))
    assert dev.snapshot.ef_remaining_time_minutes.raw_value == minutes
    assert dev.snapshot.ef_charge_remaining_time_minutes.raw_value == minutes
    assert dev.snapshot.ef_discharge_remaining_time_minutes.raw_value == minutes


def test_river3_does_not_use_river2_struct_decoder():
    dev = device("R651TEST00000001")
    dev._record(heartbeat(2, pd_payload()))
    assert dev.snapshot.ef_input_power_watts is None
    assert "ef_pkt_02_20_02" in dev.snapshot._packets


def test_fragmented_auth_probe_then_auth_and_telemetry(enc_helper):
    async def run():
        dev = device()
        dev._assembler = enc_helper
        dev._state = _State.AUTH_STATUS
        dev._write_packet = AsyncMock()
        probe = EcoFlowPacket(0x35, 0x21, 0x35, 0x89, b"\x00", version=2)
        wire = enc_helper.encode(probe)
        await dev._notify_handler(None, bytearray(wire[:9]))
        dev._write_packet.assert_not_awaited()
        assert dev._state == _State.AUTH_STATUS
        await dev._notify_handler(None, bytearray(wire[9:]))
        dev._write_packet.assert_awaited_once()
        auth = dev._write_packet.call_args.args[0]
        assert auth.version == 2
        assert auth.cmd_id == 0x86
        assert len(auth.payload) == 32
        assert dev._state == _State.AUTHENTICATING

        success = EcoFlowPacket(0x35, 0x21, 0x35, 0x86, b"\x00", version=2)
        wire = enc_helper.encode(success) + enc_helper.encode(heartbeat(2, pd_payload()))
        await dev._notify_handler(None, bytearray(wire))
        assert await dev.wait_authenticated(timeout=0.1)
        assert dev.is_authenticated
        assert dev.snapshot.ef_input_power_watts.raw_value == 160
        # Read-only collection: no echo of telemetry as device commands.
        dev._write_packet.assert_awaited_once()
    asyncio.run(run())


def test_unrecognized_f001_is_not_auth_success(enc_helper):
    async def run():
        dev = device()
        dev._state = _State.AUTHENTICATING
        dev._assembler = enc_helper
        enc = enc_helper._enc
        wire = _enc_packet_to_bytes(1, b"\xf0\x01", enc.session_key, enc.iv)
        await dev._on_data(bytearray(wire))
        assert not dev.is_authenticated
        assert dev._pkt_count == 0
        rejected = EcoFlowPacket(0x35, 0x21, 0x35, 0x86, b"\x01", version=2)
        await dev._on_data(bytearray(enc_helper.encode(rejected)))
        assert dev._state == _State.ERROR
        assert not dev._authenticated.is_set()
    asyncio.run(run())


def test_disconnect_cancels_pending_auth_fallback():
    async def run():
        dev = device()
        dev._state = _State.AUTH_STATUS
        dev._do_authenticate = AsyncMock()
        task = asyncio.create_task(dev._auth_type7_fallback())
        dev._auth_fallback_task = task
        await dev.disconnect()
        assert task.cancelled()
        assert dev._auth_fallback_task is None
        dev._do_authenticate.assert_not_awaited()
    asyncio.run(run())


@pytest.mark.parametrize("connect_ok,authenticated,packets", [
    (False, False, 0), (True, False, 0), (True, True, 0), (True, True, 1),
])
def test_collection_always_disconnects(monkeypatch, connect_ok, authenticated, packets):
    fake = Mock(
        connect=AsyncMock(return_value=connect_ok),
        wait_authenticated=AsyncMock(return_value=authenticated),
        disconnect=AsyncMock(), _pkt_count=packets,
    )
    monkeypatch.setattr("bmspy.ecoflow_bt.EcoFlowBTDevice", Mock(return_value=fake))
    result = collect_data(initialise_device(ADDRESS, SERIAL, "test-user", listen_duration=0))
    assert result is (fake.snapshot if packets else False)
    fake.disconnect.assert_awaited_once()


@pytest.fixture
def diagnostic_script():
    path = Path(__file__).resolve().parents[1] / "tools" / "ecoflow-bt.py"
    spec = importlib.util.spec_from_file_location("ecoflow_diagnostic", path)
    script = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(script)
    return script


def test_diagnostic_script_decodes_exact_source(diagnostic_script):
    lines = diagnostic_script._decode_snapshot_packet("ef_pkt_03_20_02", ems_payload().hex(), [2, 3])
    assert "EMS heartbeat" in lines[0]
    assert any("dsg_remain_time: 240 min" in line for line in lines)


def test_diagnostic_printer_preserves_empty_packet(diagnostic_script, capsys):
    dev = device()
    dev._record(EcoFlowPacket(0x35, 0x21, 0x01, 0x52, b"", version=2))
    dev._record(heartbeat(2, pd_payload()))
    diagnostic_script.print_snapshot("River 2", ADDRESS, dev.snapshot, [2, 3])
    output = capsys.readouterr().out
    assert "ef_pkt_35_01_52:" in output
    assert "(empty payload)" in output
    assert "BMSInfoField(" not in output
    assert "PD heartbeat" in output
    assert "watts_in_sum: 160 W" in output


@pytest.mark.parametrize("payload", ["BMSInfoField(info='')", "zz", "a", None])
def test_diagnostic_invalid_hex_is_reported(diagnostic_script, payload):
    lines = diagnostic_script._decode_snapshot_packet("ef_pkt_35_01_52", payload, [2, 3])
    assert "invalid hex payload" in lines[0]


@pytest.mark.parametrize("failure_stage", [None, "connect", "collect", "print", "cancel"])
def test_diagnostic_run_always_disconnects(diagnostic_script, monkeypatch, failure_stage):
    manager = Mock(
        connect_all=AsyncMock(), disconnect_all=AsyncMock(),
        add_device=Mock(return_value=device()),
    )
    collect = AsyncMock()
    printer = Mock()
    monkeypatch.setattr(diagnostic_script, "EcoFlowBTManager", Mock(return_value=manager))
    monkeypatch.setattr(diagnostic_script, "collect", collect)
    monkeypatch.setattr(diagnostic_script, "print_snapshot", printer)
    args = Namespace(scan=False, device=ADDRESS, serial_number=SERIAL,
                     debug=0, timeout=0)
    if failure_stage is not None:
        target = {
            "connect": manager.connect_all, "collect": collect,
            "print": printer, "cancel": collect,
        }[failure_stage]
        error_type = asyncio.CancelledError if failure_stage == "cancel" else RuntimeError
        target.side_effect = error_type("test failure")
        with pytest.raises(error_type):
            asyncio.run(diagnostic_script.run("test-user", args))
    else:
        asyncio.run(diagnostic_script.run("test-user", args))
        printer.assert_called_once()
    manager.disconnect_all.assert_awaited_once()
    manager.add_device.assert_called_once_with(
        address=ADDRESS, serial_number=SERIAL, encrypt_type=7,
    )


@pytest.mark.parametrize("select_device,serial_override", [
    (False, None), (True, None), (True, SERIAL),
])
@pytest.mark.parametrize("expected_type", [1, 7])
def test_diagnostic_uses_discovery_instead_of_saved_identifiers(
    diagnostic_script, monkeypatch, select_device, serial_override, expected_type,
):
    detected = SimpleNamespace(address=ADDRESS, serial_number=SERIAL, encrypt_type=expected_type, rssi=-50)
    unsupported = SimpleNamespace(address="02:00:00:00:00:02", serial_number="UNKNOWN", encrypt_type=7, rssi=-50)
    scanner = AsyncMock(return_value=[detected, unsupported])
    manager = Mock(
        connect_all=AsyncMock(), disconnect_all=AsyncMock(),
        add_device=Mock(return_value=device()),
    )
    monkeypatch.setattr(diagnostic_script, "scan_for_devices", scanner)
    monkeypatch.setattr(diagnostic_script, "EcoFlowBTManager", Mock(return_value=manager))
    monkeypatch.setattr(diagnostic_script, "collect", AsyncMock())
    monkeypatch.setattr(diagnostic_script, "print_snapshot", Mock())
    args = Namespace(scan=True, device=ADDRESS if select_device else None,
                     serial_number=serial_override, debug=0, timeout=0)
    asyncio.run(diagnostic_script.run("test-user", args))
    manager.add_device.assert_called_once_with(
        address=ADDRESS, serial_number=SERIAL, encrypt_type=expected_type,
    )
    manager.disconnect_all.assert_awaited_once()


@pytest.mark.parametrize("address,serial,scan,found,error", [
    (None, None, False, [], "Supply --device"),
    (None, SERIAL, False, [], "--serial-number requires --device"),
    (ADDRESS, None, False, [], "Supply --serial-number"),
    (ADDRESS, None, True, [], "Supply --serial-number"),
    (ADDRESS, "UNKNOWN", False, [], "Unsupported serial prefix"),
    (None, None, True, [], "No EcoFlow devices found"),
])
def test_diagnostic_missing_device_details_fail_before_connecting(
    diagnostic_script, monkeypatch, address, serial, scan, found, error,
):
    factory = Mock()
    monkeypatch.setattr(diagnostic_script, "EcoFlowBTManager", factory)
    monkeypatch.setattr(diagnostic_script, "scan_for_devices", AsyncMock(return_value=found))
    args = Namespace(scan=scan, device=address, serial_number=serial,
                     debug=0, timeout=0)
    with pytest.raises(SystemExit, match=error):
        asyncio.run(diagnostic_script.run("test-user", args))
    factory.assert_not_called()


def test_diagnostic_help_has_no_capture_key_or_encryption_options(diagnostic_script, monkeypatch, capsys):
    monkeypatch.setattr(sys, "argv", ["ecoflow-bt.py", "--help"])
    with pytest.raises(SystemExit) as exc:
        diagnostic_script.main()
    assert exc.value.code == 0
    text = capsys.readouterr().out
    assert "--keys-out" not in text
    assert "--encrypt-type" not in text
    assert "--scan" in text and "--watch" in text


@pytest.mark.parametrize("obsolete", [["--keys-out", "unused.json"], ["--encrypt-type", "7"]])
def test_diagnostic_rejects_obsolete_flags_before_connecting(diagnostic_script, monkeypatch, capsys, obsolete):
    run = Mock()
    monkeypatch.setattr(diagnostic_script, "run", run)
    monkeypatch.setattr(sys, "argv", ["ecoflow-bt.py", "--user-id", "test-user", "--scan", *obsolete])
    with pytest.raises(SystemExit) as exc:
        diagnostic_script.main()
    assert exc.value.code == 2
    assert "unrecognized arguments" in capsys.readouterr().err
    run.assert_not_called()


@pytest.mark.parametrize("serial", [SERIAL, "R651TEST00000001"])
def test_type7_live_decryption_needs_no_public_key_export(serial, monkeypatch):
    async def scenario():
        dev = device(serial)
        initial = _Type7Encryption(bytes(range(16)), bytes(range(16, 32)))
        dev._enc = initial
        session_key = bytes(range(32, 48))
        seed_data = bytes(range(18))
        derive = Mock(return_value=session_key)
        monkeypatch.setattr(dev, "_derive_session_key", derive)
        monkeypatch.setattr(dev, "_do_auth_status", AsyncMock())
        try:
            await dev._process_session_key(initial.encrypt(seed_data))
            derive.assert_called_once_with(seed_data[16:18], seed_data[:16])
            assert not hasattr(dev, "session_key")
            assert not hasattr(dev, "session_iv")
            # The active transport still owns and uses the negotiated key.
            peer = _EncFrameHelper(_Type7Encryption(session_key, initial.iv))
            packet = EcoFlowPacket(2, 0x21, 0x20, 2, b"synthetic-telemetry", version=dev._pkt_version)
            decoded = dev._assembler.decode(peer.encode(packet))
            assert len(decoded) == 1
            assert decoded[0].payload == packet.payload
        finally:
            await dev.disconnect()
    asyncio.run(scenario())