/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.pulsar.sql.presto;

import static org.apache.pulsar.common.protocol.Commands.serializeMetadataAndPayload;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.anyInt;
import static org.mockito.Mockito.anyString;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertNotNull;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.airlift.log.Logger;
import io.netty.buffer.ByteBuf;
import io.trino.spi.connector.ColumnMetadata;
import io.trino.spi.connector.ConnectorContext;
import io.trino.spi.predicate.TupleDomain;
import io.trino.testing.TestingConnectorContext;
import java.math.BigDecimal;
import java.time.LocalDate;
import java.time.LocalTime;
import java.time.ZoneId;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import javax.ws.rs.ClientErrorException;
import javax.ws.rs.core.Response;
import org.apache.bookkeeper.mledger.AsyncCallbacks;
import org.apache.bookkeeper.mledger.Entry;
import org.apache.bookkeeper.mledger.ManagedLedgerConfig;
import org.apache.bookkeeper.mledger.ManagedLedgerFactory;
import org.apache.bookkeeper.mledger.Position;
import org.apache.bookkeeper.mledger.ReadOnlyCursor;
import org.apache.bookkeeper.mledger.impl.EntryImpl;
import org.apache.bookkeeper.mledger.impl.PositionImpl;
import org.apache.bookkeeper.mledger.impl.ReadOnlyCursorImpl;
import org.apache.bookkeeper.mledger.proto.MLDataFormats;
import org.apache.bookkeeper.stats.NullStatsProvider;
import org.apache.commons.lang3.StringUtils;
import org.apache.pulsar.client.admin.Namespaces;
import org.apache.pulsar.client.admin.PulsarAdmin;
import org.apache.pulsar.client.admin.PulsarAdminException;
import org.apache.pulsar.client.admin.Schemas;
import org.apache.pulsar.client.admin.Tenants;
import org.apache.pulsar.client.admin.Topics;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.api.schema.SchemaDefinition;
import org.apache.pulsar.client.impl.schema.AvroSchema;
import org.apache.pulsar.client.impl.schema.JSONSchema;
import org.apache.pulsar.common.api.proto.MessageMetadata;
import org.apache.pulsar.common.naming.NamespaceName;
import org.apache.pulsar.common.naming.TopicName;
import org.apache.pulsar.common.partition.PartitionedTopicMetadata;
import org.apache.pulsar.common.protocol.Commands;
import org.apache.pulsar.common.schema.SchemaInfo;
import org.apache.pulsar.common.schema.SchemaType;
import org.mockito.Mockito;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.DataProvider;

public abstract class TestPulsarConnector {

    protected static final long currentTimeMicros = 1534806330000000L;

    protected PulsarConnectorConfig pulsarConnectorConfig;

    protected PulsarMetadata pulsarMetadata;

    protected PulsarAuth pulsarAuth;

    protected PulsarAdmin pulsarAdmin;

    protected Schemas schemas;

    protected PulsarSplitManager pulsarSplitManager;

    protected Map<TopicName, PulsarRecordCursor> pulsarRecordCursors = new HashMap<>();

    protected static PulsarDispatchingRowDecoderFactory dispatchingRowDecoderFactory;

    protected static final PulsarConnectorId pulsarConnectorId = new PulsarConnectorId("test-connector");

    protected static List<TopicName> topicNames;
    protected static List<TopicName> partitionedTopicNames;
    protected static Map<String, Integer> partitionedTopicsToPartitions;
    protected static Map<String, SchemaInfo> topicsToSchemas;
    protected static Map<String, Long> topicsToNumEntries;

    private static final ObjectMapper objectMapper = new ObjectMapper();

    protected static List<String> fooFieldNames = new ArrayList<>();

    protected static final NamespaceName NAMESPACE_NAME_1 = NamespaceName.get("tenant-1", "ns-1");
    protected static final NamespaceName NAMESPACE_NAME_2 = NamespaceName.get("tenant-1", "ns-2");
    protected static final NamespaceName NAMESPACE_NAME_3 = NamespaceName.get("tenant-2", "ns-1");
    protected static final NamespaceName NAMESPACE_NAME_4 = NamespaceName.get("tenant-2", "ns-2");

    protected static final TopicName TOPIC_1 = TopicName.get("persistent", NAMESPACE_NAME_1, "topic-1");
    protected static final TopicName TOPIC_2 = TopicName.get("persistent", NAMESPACE_NAME_1, "topic-2");
    protected static final TopicName TOPIC_3 = TopicName.get("persistent", NAMESPACE_NAME_2, "topic-1");
    protected static final TopicName TOPIC_4 = TopicName.get("persistent", NAMESPACE_NAME_3, "topic-1");
    protected static final TopicName TOPIC_5 = TopicName.get("persistent", NAMESPACE_NAME_4, "topic-1");
    protected static final TopicName TOPIC_6 = TopicName.get("persistent", NAMESPACE_NAME_4, "topic-2");
    protected static final TopicName NON_SCHEMA_TOPIC = TopicName.get(
        "persistent", NAMESPACE_NAME_2, "non-schema-topic");


    protected static final TopicName PARTITIONED_TOPIC_1 = TopicName.get("persistent", NAMESPACE_NAME_1,
            "partitioned-topic-1");
    protected static final TopicName PARTITIONED_TOPIC_2 = TopicName.get("persistent", NAMESPACE_NAME_1,
            "partitioned-topic-2");
    protected static final TopicName PARTITIONED_TOPIC_3 = TopicName.get("persistent", NAMESPACE_NAME_2,
            "partitioned-topic-1");
    protected static final TopicName PARTITIONED_TOPIC_4 = TopicName.get("persistent", NAMESPACE_NAME_3,
            "partitioned-topic-1");
    protected static final TopicName PARTITIONED_TOPIC_5 = TopicName.get("persistent", NAMESPACE_NAME_4,
            "partitioned-topic-1");
    protected static final TopicName PARTITIONED_TOPIC_6 = TopicName.get("persistent", NAMESPACE_NAME_4,
            "partitioned-topic-2");


    public static class Foo {
        public enum TestEnum {
            TEST_ENUM_1,
            TEST_ENUM_2,
            TEST_ENUM_3
        }

        public int field1;
        public String field2;
        public float field3;
        public double field4;
        public boolean field5;
        public long field6;
        @org.apache.avro.reflect.AvroSchema("{ \"type\": \"long\", \"logicalType\": \"timestamp-millis\" }")
        public long timestamp;
        @org.apache.avro.reflect.AvroSchema("{ \"type\": \"int\", \"logicalType\": \"time-millis\" }")
        public int time;
        @org.apache.avro.reflect.AvroSchema("{ \"type\": \"int\", \"logicalType\": \"date\" }")
        public int date;
        @org.apache.avro.reflect.AvroSchema("{ \"type\": \"bytes\", \"logicalType\": \"decimal\", \"precision\": 4, \"scale\": 2 }")
        public BigDecimal decimal;
        public TestPulsarConnector.Bar bar;
        public TestEnum field7;
    }

    public static class Bar {
        public Integer field1;
        public String field2;
        public float field3;
    }


    protected static Map<TopicName, List<PulsarColumnHandle>> topicsToColumnHandles = new HashMap<>();

    protected static Map<TopicName, PulsarSplit> splits;
    protected static Map<String, Function<Integer, Object>> fooFunctions;

    static {
        try {
            topicNames = new LinkedList<>();
            topicNames.add(TOPIC_1);
            topicNames.add(TOPIC_2);
            topicNames.add(TOPIC_3);
            topicNames.add(TOPIC_4);
            topicNames.add(TOPIC_5);
            topicNames.add(TOPIC_6);
            topicNames.add(NON_SCHEMA_TOPIC);

            partitionedTopicNames = new LinkedList<>();
            partitionedTopicNames.add(PARTITIONED_TOPIC_1);
            partitionedTopicNames.add(PARTITIONED_TOPIC_2);
            partitionedTopicNames.add(PARTITIONED_TOPIC_3);
            partitionedTopicNames.add(PARTITIONED_TOPIC_4);
            partitionedTopicNames.add(PARTITIONED_TOPIC_5);
            partitionedTopicNames.add(PARTITIONED_TOPIC_6);


            partitionedTopicsToPartitions = new HashMap<>();
            partitionedTopicsToPartitions.put(PARTITIONED_TOPIC_1.toString(), 2);
            partitionedTopicsToPartitions.put(PARTITIONED_TOPIC_2.toString(), 3);
            partitionedTopicsToPartitions.put(PARTITIONED_TOPIC_3.toString(), 4);
            partitionedTopicsToPartitions.put(PARTITIONED_TOPIC_4.toString(), 5);
            partitionedTopicsToPartitions.put(PARTITIONED_TOPIC_5.toString(), 6);
            partitionedTopicsToPartitions.put(PARTITIONED_TOPIC_6.toString(), 7);

            topicsToSchemas = new HashMap<>();
            topicsToSchemas.put(TOPIC_1.getSchemaName(), Schema.AVRO(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(TOPIC_2.getSchemaName(), Schema.AVRO(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(TOPIC_3.getSchemaName(), Schema.AVRO(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(TOPIC_4.getSchemaName(), Schema.JSON(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(TOPIC_5.getSchemaName(), Schema.JSON(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(TOPIC_6.getSchemaName(), Schema.JSON(TestPulsarMetadata.Foo.class).getSchemaInfo());

            topicsToSchemas.put(PARTITIONED_TOPIC_1.getSchemaName(), Schema.AVRO(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(PARTITIONED_TOPIC_2.getSchemaName(), Schema.AVRO(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(PARTITIONED_TOPIC_3.getSchemaName(), Schema.AVRO(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(PARTITIONED_TOPIC_4.getSchemaName(), Schema.JSON(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(PARTITIONED_TOPIC_5.getSchemaName(), Schema.JSON(TestPulsarMetadata.Foo.class).getSchemaInfo());
            topicsToSchemas.put(PARTITIONED_TOPIC_6.getSchemaName(), Schema.JSON(TestPulsarMetadata.Foo.class).getSchemaInfo());

            topicsToNumEntries = new HashMap<>();
            topicsToNumEntries.put(TOPIC_1.getSchemaName(), 1233L);
            topicsToNumEntries.put(TOPIC_2.getSchemaName(), 0L);
            topicsToNumEntries.put(TOPIC_3.getSchemaName(), 100L);
            topicsToNumEntries.put(TOPIC_4.getSchemaName(), 12345L);
            topicsToNumEntries.put(TOPIC_5.getSchemaName(), 8000L);
            topicsToNumEntries.put(TOPIC_6.getSchemaName(), 1L);

            topicsToNumEntries.put(NON_SCHEMA_TOPIC.getSchemaName(), 8000L);
            topicsToNumEntries.put(PARTITIONED_TOPIC_1.getSchemaName(), 1233L);
            topicsToNumEntries.put(PARTITIONED_TOPIC_2.getSchemaName(), 8000L);
            topicsToNumEntries.put(PARTITIONED_TOPIC_3.getSchemaName(), 100L);
            topicsToNumEntries.put(PARTITIONED_TOPIC_4.getSchemaName(), 0L);
            topicsToNumEntries.put(PARTITIONED_TOPIC_5.getSchemaName(), 800L);
            topicsToNumEntries.put(PARTITIONED_TOPIC_6.getSchemaName(), 1L);


            fooFieldNames.add("field1");
            fooFieldNames.add("field2");
            fooFieldNames.add("field3");
            fooFieldNames.add("field4");
            fooFieldNames.add("field5");
            fooFieldNames.add("field6");
            fooFieldNames.add("timestamp");
            fooFieldNames.add("time");
            fooFieldNames.add("date");
            fooFieldNames.add("bar");
            fooFieldNames.add("field7");
            fooFieldNames.add("decimal");


            ConnectorContext prestoConnectorContext = new TestingConnectorContext();
            dispatchingRowDecoderFactory = new PulsarDispatchingRowDecoderFactory(prestoConnectorContext.getTypeManager());

            topicsToColumnHandles.put(PARTITIONED_TOPIC_1, getColumnColumnHandles(PARTITIONED_TOPIC_1,topicsToSchemas.get(PARTITIONED_TOPIC_1.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(PARTITIONED_TOPIC_2, getColumnColumnHandles(PARTITIONED_TOPIC_2,topicsToSchemas.get(PARTITIONED_TOPIC_2.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(PARTITIONED_TOPIC_3, getColumnColumnHandles(PARTITIONED_TOPIC_3,topicsToSchemas.get(PARTITIONED_TOPIC_3.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(PARTITIONED_TOPIC_4, getColumnColumnHandles(PARTITIONED_TOPIC_4,topicsToSchemas.get(PARTITIONED_TOPIC_4.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(PARTITIONED_TOPIC_5, getColumnColumnHandles(PARTITIONED_TOPIC_5,topicsToSchemas.get(PARTITIONED_TOPIC_5.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(PARTITIONED_TOPIC_6, getColumnColumnHandles(PARTITIONED_TOPIC_6,topicsToSchemas.get(PARTITIONED_TOPIC_6.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));


            topicsToColumnHandles.put(TOPIC_1, getColumnColumnHandles(TOPIC_1,topicsToSchemas.get(TOPIC_1.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(TOPIC_2, getColumnColumnHandles(TOPIC_2,topicsToSchemas.get(TOPIC_2.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(TOPIC_3, getColumnColumnHandles(TOPIC_3,topicsToSchemas.get(TOPIC_3.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(TOPIC_4, getColumnColumnHandles(TOPIC_4,topicsToSchemas.get(TOPIC_4.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(TOPIC_5, getColumnColumnHandles(TOPIC_5,topicsToSchemas.get(TOPIC_5.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));
            topicsToColumnHandles.put(TOPIC_6, getColumnColumnHandles(TOPIC_6,topicsToSchemas.get(TOPIC_6.getSchemaName()), PulsarColumnHandle.HandleKeyValueType.NONE,true));


            splits = new HashMap<>();

            List<TopicName> allTopics = new LinkedList<>();
            allTopics.addAll(topicNames);
            allTopics.addAll(partitionedTopicNames);


            for (TopicName topicName : allTopics) {
                if (topicsToSchemas.containsKey(topicName.getSchemaName())) {
                    splits.put(topicName, new PulsarSplit(0, pulsarConnectorId.toString(),
                        topicName.getNamespace(), topicName.getLocalName(), topicName.getLocalName(),
                        topicsToNumEntries.get(topicName.getSchemaName()),
                        new String(topicsToSchemas.get(topicName.getSchemaName()).getSchema()),
                        topicsToSchemas.get(topicName.getSchemaName()).getType(),
                        0, topicsToNumEntries.get(topicName.getSchemaName()),
                        0, 0, TupleDomain.all(),
                            objectMapper.writeValueAsString(
                                    topicsToSchemas.get(topicName.getSchemaName()).getProperties()), null));
                }
            }

            fooFunctions = new HashMap<>();

            fooFunctions.put("field1", integer -> integer);
            fooFunctions.put("field2", String::valueOf);
            fooFunctions.put("field3", Integer::floatValue);
            fooFunctions.put("field4", Integer::doubleValue);
            fooFunctions.put("field5", integer -> integer % 2 == 0);
            fooFunctions.put("field6", Integer::longValue);
            fooFunctions.put("timestamp", integer -> System.currentTimeMillis());
            fooFunctions.put("time", integer -> {
                LocalTime now = LocalTime.now(ZoneId.systemDefault());
                return now.toSecondOfDay() * 1000;
            });
            fooFunctions.put("date", integer -> {
                LocalDate localDate = LocalDate.now();
                LocalDate epoch = LocalDate.ofEpochDay(0);
                return Math.toIntExact(ChronoUnit.DAYS.between(epoch, localDate));
            });
            fooFunctions.put("decimal", integer -> BigDecimal.valueOf(1234, 2));
            fooFunctions.put("bar.field1", integer -> integer % 3 == 0 ? null : integer + 1);
            fooFunctions.put("bar.field2", integer -> integer % 2 == 0 ? null : String.valueOf(integer + 2));
            fooFunctions.put("bar.field3", integer -> integer + 3.0f);

            fooFunctions.put("field7", integer -> Foo.TestEnum.values()[integer % Foo.TestEnum.values().length]);

        } catch (Throwable e) {
            System.out.println("Error: " + e);
            System.out.println("Stacktrace: " + Arrays.asList(e.getStackTrace()));
        }
    }


    /**
     * Parse PulsarColumnMetadata to PulsarColumnHandle Util
     * @param schemaInfo
     * @param handleKeyValueType
     * @param includeInternalColumn
     * @return
     */
    protected static List<PulsarColumnHandle> getColumnColumnHandles(TopicName topicName, SchemaInfo schemaInfo,
                                                                     PulsarColumnHandle.HandleKeyValueType handleKeyValueType, boolean includeInternalColumn) {
        List<PulsarColumnHandle> columnHandles = new ArrayList<>();
        List<ColumnMetadata> columnMetadata = mockColumnMetadata().getPulsarColumns(topicName, schemaInfo,
                includeInternalColumn, handleKeyValueType);
        columnMetadata.forEach(column -> {
            PulsarColumnMetadata pulsarColumnMetadata = (PulsarColumnMetadata) column;
            columnHandles.add(new PulsarColumnHandle(
                    pulsarConnectorId.toString(),
                    pulsarColumnMetadata.getNameWithCase(),
                    pulsarColumnMetadata.getType(),
                    pulsarColumnMetadata.isHidden(),
                    pulsarColumnMetadata.isInternal(),
                    pulsarColumnMetadata.getDecoderExtraInfo().getMapping(),
                    pulsarColumnMetadata.getDecoderExtraInfo().getDataFormat(), pulsarColumnMetadata.getDecoderExtraInfo().getFormatHint(),
                    pulsarColumnMetadata.getHandleKeyValueType()));

        });
        return columnHandles;
    }

    public static PulsarMetadata mockColumnMetadata() {
        ConnectorContext prestoConnectorContext = new TestingConnectorContext();
        PulsarConnectorConfig pulsarConnectorConfig = spy(PulsarConnectorConfig.class);
        pulsarConnectorConfig.setMaxEntryReadBatchSize(1);
        pulsarConnectorConfig.setMaxSplitEntryQueueSize(10);
        pulsarConnectorConfig.setMaxSplitMessageQueueSize(100);
        PulsarDispatchingRowDecoderFactory dispatchingRowDecoderFactory =
                new PulsarDispatchingRowDecoderFactory(prestoConnectorContext.getTypeManager());
        PulsarAuth pulsarAuth = new PulsarAuth(pulsarConnectorConfig);
        PulsarMetadata pulsarMetadata =
                new PulsarMetadata(pulsarConnectorId, pulsarConnectorConfig, dispatchingRowDecoderFactory, pulsarAuth);
        return pulsarMetadata;
    }

    public static PulsarConnectorId getPulsarConnectorId() {
        assertNotNull(pulsarConnectorId);
        return pulsarConnectorId;
    }

    private static List<Entry> getTopicEntries(String topicSchemaName) {
        List<Entry> entries = new LinkedList<>();

        long count = topicsToNumEntries.get(topicSchemaName);
        for (int i=0 ; i < count; i++) {

            Foo foo = new Foo();
            foo.field1 = (int) count;
            foo.field2 = String.valueOf(count);
            foo.field3 = count;
            foo.field4 = count;
            foo.field5 = count % 2 == 0;
            foo.field6 = count;
            foo.timestamp = System.currentTimeMillis();

            LocalTime now = LocalTime.now(ZoneId.systemDefault());
            foo.time = now.toSecondOfDay() * 1000;

            LocalDate localDate = LocalDate.now();
            LocalDate epoch = LocalDate.ofEpochDay(0);
            foo.date = Math.toIntExact(ChronoUnit.DAYS.between(epoch, localDate));
            foo.decimal= BigDecimal.valueOf(count, 2);

            MessageMetadata messageMetadata = new MessageMetadata()
                    .setProducerName("test-producer").setSequenceId(i)
                    .setPublishTime(currentTimeMicros / 1000 + i);

            Schema schema = topicsToSchemas.get(topicSchemaName).getType() == SchemaType.AVRO ? AvroSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).build()) : JSONSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).build());

            ByteBuf payload = io.netty.buffer.Unpooled
                    .copiedBuffer(schema.encode(foo));

            ByteBuf byteBuf = serializeMetadataAndPayload(
                    Commands.ChecksumType.Crc32c, messageMetadata, payload);

            Entry entry = EntryImpl.create(0, i, byteBuf);
            log.info("create entry: %s", entry.getEntryId());
            entries.add(entry);
        }
        return entries;
    }

    public long completedBytes = 0L;

    private static final Logger log = Logger.get(TestPulsarConnector.class);

    protected static List<String> getNamespace(String tenant) {
        return topicNames.stream()
            .filter(topicName -> topicName.getTenant().equals(tenant))
            .map(TopicName::getNamespace)
            .distinct()
            .collect(Collectors.toCollection(LinkedList::new));
    }

    protected static List<String> getTopics(String ns) {
        List<String> topics = new ArrayList<>(topicNames.stream()
            .filter(topicName -> topicName.getNamespace().equals(ns))
            .map(TopicName::toString).collect(Collectors.toList()));
        partitionedTopicNames.stream().filter(topicName -> topicName.getNamespace().equals(ns)).forEach(topicName -> {
            for (Integer i = 0; i < partitionedTopicsToPartitions.get(topicName.toString()); i++) {
                topics.add(TopicName.get(topicName + "-partition-" + i).toString());
            }
        });
        return topics;
    }

    protected static List<String> getPartitionedTopics(String ns) {
        return partitionedTopicNames.stream()
            .filter(topicName -> topicName.getNamespace().equals(ns))
            .map(TopicName::toString)
            .collect(Collectors.toList());
    }

    @BeforeMethod
    public void setup() throws Exception {
        this.pulsarConnectorConfig = spy(PulsarConnectorConfig.class);
        this.pulsarConnectorConfig.setMaxEntryReadBatchSize(1);
        this.pulsarConnectorConfig.setMaxSplitEntryQueueSize(10);
        this.pulsarConnectorConfig.setMaxSplitMessageQueueSize(100);

        Tenants tenants = mock(Tenants.class);
        doReturn(new LinkedList<>(topicNames.stream()
            .map(TopicName::getTenant)
            .collect(Collectors.toSet()))).when(tenants).getTenants();

        Namespaces namespaces = mock(Namespaces.class);

        when(namespaces.getNamespaces(anyString())).thenAnswer(new Answer<List<String>>() {
            @Override
            public List<String> answer(InvocationOnMock invocation) throws Throwable {
                Object[] args = invocation.getArguments();
                String tenant = (String) args[0];
                List<String> ns = getNamespace(tenant);
                if (ns.isEmpty()) {
                    ClientErrorException cee = new ClientErrorException(Response.status(404).build());
                    throw new PulsarAdminException(cee, cee.getMessage(), cee.getResponse().getStatus());
                }
                return ns;
            }
        });

        Topics topics = mock(Topics.class);
        when(topics.getList(anyString(), any())).thenAnswer(new Answer<List<String>>() {
            @Override
            public List<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
                Object[] args = invocationOnMock.getArguments();
                String ns = (String) args[0];
                List<String> topics = getTopics(ns);
                if (topics.isEmpty()) {
                    ClientErrorException cee = new ClientErrorException(Response.status(404).build());
                    throw new PulsarAdminException(cee, cee.getMessage(), cee.getResponse().getStatus());
                }
                return topics;
            }
        });

        when(topics.getPartitionedTopicList(anyString())).thenAnswer(new Answer<List<String>>() {
            @Override
            public List<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
                Object[] args = invocationOnMock.getArguments();
                String ns = (String) args[0];
                List<String> topics = getPartitionedTopics(ns);
                if (topics.isEmpty()) {
                    ClientErrorException cee = new ClientErrorException(Response.status(404).build());
                    throw new PulsarAdminException(cee, cee.getMessage(), cee.getResponse().getStatus());
                }
                return topics;
            }
        });

        when(topics.getPartitionedTopicMetadata(anyString())).thenAnswer(new Answer<PartitionedTopicMetadata>() {
            @Override
            public PartitionedTopicMetadata answer(InvocationOnMock invocationOnMock) throws Throwable {
                Object[] args = invocationOnMock.getArguments();
                String topic = (String) args[0];
                int partitions = partitionedTopicsToPartitions.get(topic) == null
                        ? 0 : partitionedTopicsToPartitions.get(topic);
                return new PartitionedTopicMetadata(partitions);
            }
        });

        schemas = mock(Schemas.class);
        when(schemas.getSchemaInfo(anyString())).thenAnswer(new Answer<SchemaInfo>() {
            @Override
            public SchemaInfo answer(InvocationOnMock invocationOnMock) throws Throwable {
                Object[] args = invocationOnMock.getArguments();
                String topic = (String) args[0];
                if (topicsToSchemas.get(topic) != null) {
                    return topicsToSchemas.get(topic);
                } else {
                    ClientErrorException cee = new ClientErrorException(Response.status(404).build());
                    throw new PulsarAdminException(cee, cee.getMessage(), cee.getResponse().getStatus());
                }
            }
        });

        pulsarAdmin = mock(PulsarAdmin.class);
        doReturn(tenants).when(pulsarAdmin).tenants();
        doReturn(namespaces).when(pulsarAdmin).namespaces();
        doReturn(topics).when(pulsarAdmin).topics();
        doReturn(schemas).when(pulsarAdmin).schemas();
        doReturn(pulsarAdmin).when(this.pulsarConnectorConfig).getPulsarAdmin();

        this.pulsarAuth = mock(PulsarAuth.class);

        this.pulsarMetadata =
                new PulsarMetadata(pulsarConnectorId, this.pulsarConnectorConfig, dispatchingRowDecoderFactory,
                        this.pulsarAuth);
        this.pulsarSplitManager = Mockito.spy(new PulsarSplitManager(pulsarConnectorId, this.pulsarConnectorConfig));

        ManagedLedgerFactory managedLedgerFactory = mock(ManagedLedgerFactory.class);
        when(managedLedgerFactory.openReadOnlyCursor(any(), any(), any())).then(new Answer<ReadOnlyCursor>() {

            private Map<String, Integer> positions = new HashMap<>();

            private int count = 0;
            @Override
            public ReadOnlyCursor answer(InvocationOnMock invocationOnMock) throws Throwable {
                Object[] args = invocationOnMock.getArguments();
                String topic = (String) args[0];
                PositionImpl positionImpl = (PositionImpl) args[1];

                int position = positionImpl.getEntryId() == -1 ? 0 : (int) positionImpl.getEntryId();

                positions.put(topic, position);
                String schemaName = TopicName.get(
                        TopicName.get(
                                topic.replaceAll("/persistent", ""))
                                .getPartitionedTopicName()).getSchemaName();
                long entries = topicsToNumEntries.get(schemaName);


                ReadOnlyCursorImpl readOnlyCursor = mock(ReadOnlyCursorImpl.class);
                doReturn(entries).when(readOnlyCursor).getNumberOfEntries();

                doAnswer(new Answer<Void>() {
                    @Override
                    public Void answer(InvocationOnMock invocation) throws Throwable {
                        Object[] args = invocation.getArguments();
                        Integer skipEntries = (Integer) args[0];
                        positions.put(topic, positions.get(topic) + skipEntries);
                        return null;
                    }
                }).when(readOnlyCursor).skipEntries(anyInt());

                when(readOnlyCursor.getReadPosition()).thenAnswer(new Answer<PositionImpl>() {
                    @Override
                    public PositionImpl answer(InvocationOnMock invocationOnMock) throws Throwable {
                        return PositionImpl.get(0, positions.get(topic));
                    }
                });

                doAnswer(new Answer() {
                    @Override
                    public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
                        Object[] args = invocationOnMock.getArguments();
                        Integer readEntries = (Integer) args[0];
                        AsyncCallbacks.ReadEntriesCallback callback = (AsyncCallbacks.ReadEntriesCallback) args[2];
                        Object ctx = args[3];

                        new Thread(new Runnable() {
                            @Override
                            public void run() {
                                List <Entry> entries = new LinkedList<>();
                                for (int i = 0; i < readEntries; i++) {

                                    TestPulsarConnector.Bar bar = new TestPulsarConnector.Bar();
                                    bar.field1 = fooFunctions.get("bar.field1").apply(count) == null ? null : (int) fooFunctions.get("bar.field1").apply(count);
                                    bar.field2 = fooFunctions.get("bar.field2").apply(count) == null ? null : (String) fooFunctions.get("bar.field2").apply(count);
                                    bar.field3 = (float) fooFunctions.get("bar.field3").apply(count);


                                    Foo foo = new Foo();
                                    foo.field1 = (int) fooFunctions.get("field1").apply(count);
                                    foo.field2 = (String) fooFunctions.get("field2").apply(count);
                                    foo.field3 = (float) fooFunctions.get("field3").apply(count);
                                    foo.field4 = (double) fooFunctions.get("field4").apply(count);
                                    foo.field5 = (boolean) fooFunctions.get("field5").apply(count);
                                    foo.field6 = (long) fooFunctions.get("field6").apply(count);
                                    foo.timestamp = (long) fooFunctions.get("timestamp").apply(count);
                                    foo.time = (int) fooFunctions.get("time").apply(count);
                                    foo.date = (int) fooFunctions.get("date").apply(count);
                                    foo.decimal = (BigDecimal) fooFunctions.get("decimal").apply(count);
                                    foo.bar = bar;
                                    foo.field7 = (Foo.TestEnum) fooFunctions.get("field7").apply(count);

                                    MessageMetadata messageMetadata = new MessageMetadata()
                                            .setProducerName("test-producer").setSequenceId(positions.get(topic))
                                            .setPublishTime(System.currentTimeMillis());

                                    Schema schema = topicsToSchemas.get(schemaName).getType() == SchemaType.AVRO ? AvroSchema.of(Foo.class) : JSONSchema.of(Foo.class);

                                    ByteBuf payload = io.netty.buffer.Unpooled
                                            .copiedBuffer(schema.encode(foo));

                                    ByteBuf byteBuf = serializeMetadataAndPayload(
                                            Commands.ChecksumType.Crc32c, messageMetadata, payload);

                                    completedBytes += byteBuf.readableBytes();

                                    entries.add(EntryImpl.create(0, positions.get(topic), byteBuf));
                                    positions.put(topic, positions.get(topic) + 1);
                                    count++;
                                }

                                callback.readEntriesComplete(entries, ctx);
                            }
                        }).start();

                        return null;
                    }
                }).when(readOnlyCursor).asyncReadEntries(anyInt(), anyLong(), any(), any(), any());

                when(readOnlyCursor.hasMoreEntries()).thenAnswer(new Answer<Boolean>() {
                    @Override
                    public Boolean answer(InvocationOnMock invocationOnMock) throws Throwable {
                        return positions.get(topic) < entries;
                    }
                });

                when(readOnlyCursor.findNewestMatching(any(), any())).then(new Answer<Position>() {
                    @Override
                    public Position answer(InvocationOnMock invocationOnMock) throws Throwable {
                        Object[] args = invocationOnMock.getArguments();
                        Predicate<Entry> predicate
                                = (Predicate<Entry>) args[1];

                        String schemaName = TopicName.get(
                                TopicName.get(
                                        topic.replaceAll("/persistent", ""))
                                        .getPartitionedTopicName()).getSchemaName();
                        List<Entry> entries = getTopicEntries(schemaName);

                        Integer target = null;
                        for (int i=entries.size() - 1; i >= 0; i--) {
                            Entry entry = entries.get(i);
                            if (predicate.test(entry)) {
                                target = i;
                                break;
                            }
                        }

                        return target == null ? null : new PositionImpl(0, target);
                    }
                });

                when(readOnlyCursor.getNumberOfEntries(any())).then(new Answer<Long>() {
                    @Override
                    public Long answer(InvocationOnMock invocationOnMock) throws Throwable {
                        Object[] args = invocationOnMock.getArguments();
                        com.google.common.collect.Range<PositionImpl>  range
                                = (com.google.common.collect.Range<PositionImpl> ) args[0];

                        return (range.upperEndpoint().getEntryId() + 1) - range.lowerEndpoint().getEntryId();
                    }
                });

                when(readOnlyCursor.getCurrentLedgerInfo()).thenReturn(MLDataFormats.ManagedLedgerInfo.LedgerInfo.newBuilder().setLedgerId(0).build());

                return readOnlyCursor;
            }
        });

        PulsarConnectorCache.instance = mock(PulsarConnectorCache.class);
        when(PulsarConnectorCache.instance.getManagedLedgerFactory()).thenReturn(managedLedgerFactory);

        for (Map.Entry<TopicName, PulsarSplit> split : splits.entrySet()) {
            PulsarRecordCursor pulsarRecordCursor = new PulsarRecordCursor(
                    topicsToColumnHandles.get(split.getKey()), split.getValue(),
                    pulsarConnectorConfig, managedLedgerFactory, new ManagedLedgerConfig(),
                    new PulsarConnectorMetricsTracker(new NullStatsProvider()),dispatchingRowDecoderFactory);
            this.pulsarRecordCursors.put(split.getKey(), pulsarRecordCursor);
        }
    }

    @AfterMethod(alwaysRun = true)
    public void cleanup() {
        completedBytes = 0L;
    }

    @DataProvider(name = "rewriteNamespaceDelimiter")
    public static Object[][] serviceUrls() {
        return new Object[][] {
                { "|" }, { null }
        };
    }

    protected void updateRewriteNamespaceDelimiterIfNeeded(String delimiter) {
        if (StringUtils.isNotBlank(delimiter)) {
            pulsarConnectorConfig.setNamespaceDelimiterRewriteEnable(true);
            pulsarConnectorConfig.setRewriteNamespaceDelimiter(delimiter);
        } else {
            pulsarConnectorConfig.setNamespaceDelimiterRewriteEnable(false);
        }
    }
}
