from unittest.mock import MagicMock

from awe.mq.mq_producer import MQProducer, MQProducerOptions
from awe.types.common import Payload, ResponseMessage


def _build_payload(task_id: str):
    return Payload(
        taskId=task_id,
        userId='user-1',
        clientId='client-1',
        workflowCode='wf-1',
        messageTag='APP',
    )


def _build_progress_message(message_type='progress'):
    return ResponseMessage({
        'type': 'PROGRESS',
        'debounced': True,
        'response': {
            'progress': 10,
            'progressData': {
                'type': message_type,
            },
        },
    })


def test_send_progress_message_throttles_same_task(monkeypatch):
    options = MQProducerOptions(progress_throttle_seconds=1)
    monkeypatch.setattr(MQProducer, 'get_producer', lambda self, options: MagicMock())
    producer = MQProducer(options)
    producer.publish_message = MagicMock()

    timestamps = iter([0.0, 0.2, 1.2])
    monkeypatch.setattr('awe.mq.mq_producer.time.monotonic', lambda: next(timestamps))

    payload = _build_payload('task-1')
    message = _build_progress_message()

    producer.send_progress_message(message, payload)
    producer.send_progress_message(message, payload)
    producer.send_progress_message(message, payload)

    assert producer.publish_message.call_count == 2


def test_send_progress_message_does_not_cross_throttle_tasks(monkeypatch):
    options = MQProducerOptions(progress_throttle_seconds=1)
    monkeypatch.setattr(MQProducer, 'get_producer', lambda self, options: MagicMock())
    producer = MQProducer(options)
    producer.publish_message = MagicMock()

    timestamps = iter([0.0, 0.1])
    monkeypatch.setattr('awe.mq.mq_producer.time.monotonic', lambda: next(timestamps))

    producer.send_progress_message(_build_progress_message(), _build_payload('task-1'))
    producer.send_progress_message(_build_progress_message(), _build_payload('task-2'))

    assert producer.publish_message.call_count == 2


def test_send_complete_message_clears_progress_state(monkeypatch):
    options = MQProducerOptions(progress_throttle_seconds=1)
    monkeypatch.setattr(MQProducer, 'get_producer', lambda self, options: MagicMock())
    producer = MQProducer(options)
    producer.publish_message = MagicMock()

    timestamps = iter([0.0, 0.1])
    monkeypatch.setattr('awe.mq.mq_producer.time.monotonic', lambda: next(timestamps))

    payload = _build_payload('task-1')
    producer.send_progress_message(_build_progress_message(), payload)
    producer.send_complete_message(ResponseMessage({'type': 'COMPLETED', 'response': {'data': {}}}), payload)
    producer.send_progress_message(_build_progress_message(), payload)

    assert producer.publish_message.call_count == 3
