from queue import Empty, Queue
import time
from uuid import UUID, uuid4

from dynalab_core.derive import DeriveRegistry
from dynalab_core.protocols.packets.data import ValueDescriptor
from dynalab_core.protocols.packets.handshake import SignalDescriptor


PROCESSING_FUNCTION = """
from dynalab_core.protocols.packets.data import ValueDescriptor
from dynalab_core.protocols.packets.handshake import SignalDescriptor

def add(left: ValueDescriptor, right: ValueDescriptor, output: SignalDescriptor) -> ValueDescriptor:
    return ValueDescriptor(
        signal_id=output.id,
        value=left.value + right.value,
        timestamp=left.timestamp,
    )
"""


def test_derive_registry_routes_with_live_values() -> None:
    output_queue: Queue[ValueDescriptor] = Queue()
    left_signal = SignalDescriptor(id=uuid4(), name="Left", type="number")
    right_signal = SignalDescriptor(id=uuid4(), name="Right", type="number")
    output_signal = SignalDescriptor(id=uuid4(), name="Total", type="number")
    live_values: dict[UUID, ValueDescriptor] = {}
    registry = DeriveRegistry(output_queue, live_values.get)

    try:
        unit = registry.register(
            PROCESSING_FUNCTION, [left_signal, right_signal], output_signal
        )
        right_value = ValueDescriptor(
            signal_id=right_signal.id, value=2.0, timestamp=1
        )
        live_values[right_signal.id] = right_value
        left_value = ValueDescriptor(signal_id=left_signal.id, value=3.0, timestamp=2)

        registry.put_data(left_value)

        deadline = time.monotonic() + 1.0
        while True:
            try:
                result = output_queue.get_nowait()
                break
            except Empty:
                if time.monotonic() >= deadline:
                    raise AssertionError("Derive unit did not produce a value")
                time.sleep(0.01)

        registry.unregister(unit.uuid())
        assert registry.get(unit.uuid()) is None
    finally:
        registry.stop()

    assert result.signal_id == output_signal.id
    assert result.value == 5.0
    assert result.timestamp == left_value.timestamp

def test_derive_registry_ignores_unknown_unit_on_unregister() -> None:
    registry = DeriveRegistry(Queue(), lambda signal_id: None)

    try:
        registry.unregister(uuid4())
    finally:
        registry.stop()