from __future__ import annotations

import pytest
from fivetran_connector_sdk import operation_stream
from fivetran_connector_sdk.constants import UPSERT_TYPE
from fivetran_connector_sdk.protos import connector_sdk_pb2

from fivetran.batched_stream import install_batched_operation_stream_support


def test_batched_operation_stream_drains_record_batch() -> None:
    install_batched_operation_stream_support()
    stream = operation_stream._OperationStream()
    records = [_record(1), _record(2)]

    stream.add_records(records)
    stream.mark_done()

    response = next(stream)
    assert list(response.records.records) == records
    with pytest.raises(StopIteration):
        next(stream)


def test_batched_operation_stream_preserves_regular_record_order() -> None:
    install_batched_operation_stream_support()
    stream = operation_stream._OperationStream()
    records = [_record(1), _record(2), _record(3)]

    stream.add_record(records[0])
    stream.add_records(records[1:])
    stream.mark_done()

    first_response = next(stream)
    second_response = next(stream)

    assert list(first_response.records.records) == [records[0]]
    assert list(second_response.records.records) == records[1:]
    with pytest.raises(StopIteration):
        next(stream)


def test_batched_operation_stream_emits_record_batch_before_checkpoint() -> None:
    install_batched_operation_stream_support()
    stream = operation_stream._OperationStream()
    records = [_record(1), _record(2)]
    checkpoint = connector_sdk_pb2.Checkpoint(state_json='{"checkpoint": true}')

    stream.add_records(records)
    stream._queue.put(checkpoint)
    stream.mark_done()

    first_response = next(stream)
    second_response = next(stream)

    assert list(first_response.records.records) == records
    assert len(second_response) == 1
    assert second_response[0].checkpoint == checkpoint
    with pytest.raises(StopIteration):
        next(stream)


def _record(value: int) -> connector_sdk_pb2.Record:
    record = connector_sdk_pb2.Record(table_name="test_table", type=UPSERT_TYPE)
    record.data["id"].long = value
    return record