# Path: cb_bin_client.py
# def decode_collection_id(key: bytes) -> Tuple[int, bytes]:
#     # A leb128 varint encodes the CID
#     data = array.array('B')
#     data.frombytes(key)
#     cid = data[0] & 0x7f
#     end = 1
#     if (data[0] & 0x80) == 0x80:
#         shift = 7
#         for end in range(1, len(data)):
#             cid |= ((data[end] & 0x7f) << shift)
#             if (data[end] & 0x80) == 0:
#                 break
#             shift = shift + 7
# 
#         end = end + 1
#         if end == len(data):
#             #  We should of stopped for a stop byte, not the end of the buffer
#             raise ValueError("encoded key did not contain a stop byte")
#     return cid, key[end:]
# 
# def encode_collection_id(cid: int) -> bytes:
#     output = array.array('B', [0])
#     while cid > 0:
#         byte = cid & 0xFF
#         cid >>= 7
#         # CID has more bits
#         if cid > 0:
#             # Set the 'continue' bit of this byte
#             byte |= 0x80
#             output[-1] = byte
#             output.append(0)
#         else:
#             output[-1] = byte
#     return output.tobytes()


import glob
import json
import logging
import os
import queue
import re
import struct
import threading
import couchstore  # pylint: disable=import-error
import couchbaseConstants
import pump
from collections import defaultdict
from typing import Any, Dict, List, Optional, Tuple
from cb_bin_client import decode_collection_id, encode_collection_id

    @staticmethod
    def run(self):
        destination_vbucket_state = getattr(self.opts, 'destination_vbucket_state', 'active')

        vbucket_states = self.source_node.get('vbucket_states', {})

        while not self.ctl['stop']:
            batch, future = self.pull_next_batch()  # type: pump.Batch, pump.SinkBatchFuture
            if not batch:
                return self.future_done(future, 0)

            vbuckets = batch.group_by_vbucket_id(SFD_VBUCKETS, self.rehash)
            for vbucket_id, msgs in vbuckets.items():
                checkpoint_id = 0
                max_deleted_seqno = 0

                rv, store, store_path = self.open_store(vbucket_id)
                if rv != 0:
                    return self.future_done(future, rv)

                bulk_keys = []
                bulk_vals = []

                for i, msg in enumerate(msgs):
                    cmd, _vbucket_id, key, flg, exp, cas, meta, val, seqno, dtype, nmeta, conf_res = msg
                    if self.skip(key, vbucket_id):
                        continue

                    # TODO: add default collection to all keys in CC this should change to have the correct collection

# Path: cb_bin_client.py
# def decode_collection_id(key: bytes) -> Tuple[int, bytes]:
#     # A leb128 varint encodes the CID
#     data = array.array('B')
#     data.frombytes(key)
#     cid = data[0] & 0x7f
#     end = 1
#     if (data[0] & 0x80) == 0x80:
#         shift = 7
#         for end in range(1, len(data)):
#             cid |= ((data[end] & 0x7f) << shift)
#             if (data[end] & 0x80) == 0:
#                 break
#             shift = shift + 7
# 
#         end = end + 1
#         if end == len(data):
#             #  We should of stopped for a stop byte, not the end of the buffer
#             raise ValueError("encoded key did not contain a stop byte")
#     return cid, key[end:]
# 
# def encode_collection_id(cid: int) -> bytes:
#     output = array.array('B', [0])
#     while cid > 0:
#         byte = cid & 0xFF
#         cid >>= 7
#         # CID has more bits
#         if cid > 0:
#             # Set the 'continue' bit of this byte
#             byte |= 0x80
#             output[-1] = byte
#             output.append(0)
#         else:
#             output[-1] = byte
#     return output.tobytes()


# Path: pump_sfd.py
import glob
import json
import logging
import os
import queue
import re
import struct
import threading
import couchstore  # pylint: disable=import-error
import couchbaseConstants
import pump
from collections import defaultdict
from typing import Any, Dict, List, Optional, Tuple
from cb_bin_client import decode_collection_id, encode_collection_id


    @staticmethod
    def run(self):
        destination_vbucket_state = getattr(self.opts, 'destination_vbucket_state', 'active')

        vbucket_states = self.source_node.get('vbucket_states', {})

        while not self.ctl['stop']:
            batch, future = self.pull_next_batch()  # type: pump.Batch, pump.SinkBatchFuture
            if not batch:
                return self.future_done(future, 0)

            vbuckets = batch.group_by_vbucket_id(SFD_VBUCKETS, self.rehash)
            for vbucket_id, msgs in vbuckets.items():
                checkpoint_id = 0
                max_deleted_seqno = 0

                rv, store, store_path = self.open_store(vbucket_id)
                if rv != 0:
                    return self.future_done(future, rv)

                bulk_keys = []
                bulk_vals = []

                for i, msg in enumerate(msgs):
                    cmd, _vbucket_id, key, flg, exp, cas, meta, val, seqno, dtype, nmeta, conf_res = msg
                    if self.skip(key, vbucket_id):
                        continue

                    # TODO: add default collection to all keys in CC this should change to have the correct collection

                    key = encode_collection_id(0) + key