from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from fivetran_connector_sdk import operation_stream
from fivetran_connector_sdk.protos import connector_sdk_pb2 as sdk_pb2


@dataclass(frozen=True)
class _RecordBatch:
    records: list[Any]


def install_batched_operation_stream_support() -> None:
    stream_type: Any = operation_stream._OperationStream
    if getattr(stream_type, "_batch_support_installed", False):
        return

    original_init = stream_type.__init__

    def __init__(self: Any, *args: Any, **kwargs: Any) -> None:
        original_init(self, *args, **kwargs)
        _ensure_batch_slots(self)

    def add_records(self: Any, records: list[Any]) -> None:
        if not records:
            return
        _ensure_batch_slots(self)
        self._batch_slots.acquire()
        try:
            self._queue.put(_RecordBatch(records))
        except BaseException:
            self._batch_slots.release()
            raise

    def _build_next_batch(self: Any) -> Any:
        while self._buffer_record_count < operation_stream.MAX_RECORDS_IN_BATCH and (
            self._buffer_size_bytes < operation_stream.MAX_BATCH_SIZE_IN_BYTES
        ):
            operation = _next_operation(self)

            if operation is self._sentinel:
                self._is_done = True
                if self._buffer:
                    return self._flush_buffer()
                raise StopIteration

            if isinstance(operation, sdk_pb2.Checkpoint):
                return self._flush_buffer_on_checkpoint(operation)

            if isinstance(operation, _RecordBatch):
                if not operation.records:
                    continue
                if self._buffer:
                    # The original SDK can already have regular records in its
                    # internal buffer. Flush those first so a later batch never
                    # overtakes earlier standard Operations.* records.
                    _prepend_operation(self, operation)
                    return self._flush_buffer()
                self._batch_slots.release()
                return sdk_pb2.UpdateResponse(records=sdk_pb2.Records(records=operation.records))

            self._buffer_record_count += 1
            self._buffer_size_bytes += operation.ByteSize()
            self._buffer.append(operation)

        return self._flush_buffer()

    stream_type.__init__ = __init__
    stream_type.add_records = add_records
    stream_type._build_next_batch = _build_next_batch
    stream_type._batch_support_installed = True


def _next_operation(stream: Any) -> Any:
    pending = getattr(stream, "_pending_operations", None)
    if pending:
        return pending.pop()
    return stream._queue.get()


def _prepend_operation(stream: Any, operation: Any) -> None:
    pending = getattr(stream, "_pending_operations", None)
    if pending is None:
        pending = []
        stream._pending_operations = pending
    pending.append(operation)


def _ensure_batch_slots(stream: Any) -> None:
    if hasattr(stream, "_batch_slots"):
        return
    # A SDK queue item can now hold up to MAX_RECORDS_IN_BATCH records. Keep
    # only one producer-side batch outstanding so this extension does not
    # multiply the SDK queue's original memory envelope.
    stream._batch_slots = operation_stream.threading.BoundedSemaphore(1)


__all__ = ["install_batched_operation_stream_support"]