/*
 * Copyright 2023 CeresDB Project Authors. Licensed under Apache-2.0.
 */
package io.ceresdb;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import io.ceresdb.models.Row;
import io.ceresdb.proto.internal.Common;
import io.ceresdb.proto.internal.Storage;

import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.Float4Vector;
import org.apache.arrow.vector.Float8Vector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.SmallIntVector;
import org.apache.arrow.vector.TimeStampMilliVector;
import org.apache.arrow.vector.TinyIntVector;
import org.apache.arrow.vector.UInt1Vector;
import org.apache.arrow.vector.UInt2Vector;
import org.apache.arrow.vector.UInt4Vector;
import org.apache.arrow.vector.UInt8Vector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
import org.apache.arrow.vector.types.pojo.Field;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.Mockito;
import org.mockito.runners.MockitoJUnitRunner;

import io.ceresdb.common.Endpoint;
import io.ceresdb.common.util.internal.ThrowUtil;
import io.ceresdb.errors.IteratorException;
import io.ceresdb.errors.StreamException;
import io.ceresdb.models.Err;
import io.ceresdb.models.SqlQueryOk;
import io.ceresdb.models.SqlQueryRequest;
import io.ceresdb.models.Result;
import io.ceresdb.options.QueryOptions;
import io.ceresdb.rpc.Context;
import io.ceresdb.rpc.Observer;
import io.ceresdb.util.Utils;

import com.google.protobuf.ByteStringHelper;

@RunWith(value = MockitoJUnitRunner.class)
public class QueryClientTest {

    private QueryClient  queryClient;
    @Mock
    private RouterClient routerClient;

    @Before
    public void before() {
        final QueryOptions queryOpts = new QueryOptions();
        queryOpts.setAsyncPool(ForkJoinPool.commonPool());
        queryOpts.setRouterClient(this.routerClient);
        queryOpts.setDatabase("public");

        this.queryClient = new QueryClient();
        this.queryClient.init(queryOpts);
    }

    @After
    public void after() {
        this.queryClient.shutdownGracefully();
        this.routerClient.shutdownGracefully();
    }

    @Test
    public void queryOkNoRouteTest() throws ExecutionException, InterruptedException, IOException {
        final Storage.SqlQueryResponse resp = mockSimpleQueryResponse(1, false);
        final Endpoint ep = Endpoint.of("127.0.0.1", 8081);

        Mockito.when(this.routerClient.invoke(Mockito.eq(ep), Mockito.any(), Mockito.any())) //
                .thenReturn(Utils.completedCf(resp));
        Mockito.when(this.routerClient.routeFor(Mockito.any(), Mockito.any())) //
                .thenReturn(Utils.completedCf(new HashMap<>()));
        Mockito.when(this.routerClient.clusterRoute()) //
                .thenReturn(Route.of(ep));

        final SqlQueryRequest req = SqlQueryRequest.newBuilder().forTables("query_test_table") //
                .sql("select number from query_test_table") //
                .build();
        final CompletableFuture<Result<SqlQueryOk, Err>> f = this.queryClient.sqlQuery(req, Context.newDefault());

        final Result<SqlQueryOk, Err> r = f.get();

        Assert.assertTrue(r.isOk());

        final SqlQueryOk queryOk = r.getOk();

        final Stream<String> strs = queryOk.map(Row::toString);

        Assert.assertEquals(
                Collections.singletonList("t1:Value{type=String,value=tvtest}|f1:Value{type=Int32,value=123}"),
                strs.collect(Collectors.toList()));
    }

    @Test
    public void queryOkByValidRouteTest() throws ExecutionException, InterruptedException, IOException {
        final Storage.SqlQueryResponse resp = mockSimpleQueryResponse(1, false);
        final Endpoint ep = Endpoint.of("127.0.0.1", 8081);

        Mockito.when(this.routerClient.invoke(Mockito.eq(ep), Mockito.any(), Mockito.any())) //
                .thenReturn(Utils.completedCf(resp));
        Mockito.when(
                this.routerClient.routeFor(Mockito.any(), Mockito.eq(Collections.singletonList("query_test_table")))) //
                .thenReturn(Utils.completedCf(new HashMap<String, Route>() {
                    private static final long serialVersionUID = -6260265905617276356L;

                    {
                        put("query_test_table", Route.of("query_test_table", ep));
                    }
                }));

        final SqlQueryRequest req = SqlQueryRequest.newBuilder().forTables("query_test_table") //
                .sql("select number from query_test_table") //
                .build();
        final CompletableFuture<Result<SqlQueryOk, Err>> f = this.queryClient.sqlQuery(req, Context.newDefault());

        final Result<SqlQueryOk, Err> r = f.get();

        Assert.assertTrue(r.isOk());

        final SqlQueryOk queryOk = r.getOk();

        final Stream<String> strs = queryOk.map(Row::toString);

        Assert.assertEquals(
                Collections.singletonList("t1:Value{type=String,value=tvtest}|f1:Value{type=Int32,value=123}"),
                strs.collect(Collectors.toList()));
    }

    @Test
    public void queryFailedTest() throws ExecutionException, InterruptedException, IOException {
        final Common.ResponseHeader header = Common.ResponseHeader.newBuilder() //
                .setCode(500) // failed code
                .build();
        final Storage.SqlQueryResponse resp = Storage.SqlQueryResponse.newBuilder().setHeader(header) //
                .build();
        final Endpoint ep = Endpoint.of("127.0.0.1", 8081);

        Mockito.when(this.routerClient.invoke(Mockito.eq(ep), Mockito.any(), Mockito.any())) //
                .thenReturn(Utils.completedCf(resp));
        Mockito.when(
                this.routerClient.routeFor(Mockito.any(), Mockito.eq(Collections.singletonList("query_test_table")))) //
                .thenReturn(Utils.completedCf(new HashMap<String, Route>() {
                    private static final long serialVersionUID = 4517371149948738282L;

                    {
                        put("query_test_table", Route.of("query_test_table", ep));
                    }
                }));
        Mockito.when(this.routerClient.routeRefreshFor(Mockito.any(), Mockito.any()))
                .thenReturn(Utils.completedCf(new HashMap<String, Route>() {
                    private static final long serialVersionUID = 2347114952231996366L;

                    {
                        put("query_test_table", Route.of("query_test_table", ep));
                    }
                }));

        final SqlQueryRequest req = SqlQueryRequest.newBuilder().forTables("query_test_table") //
                .sql("select number from query_test_table") //
                .build();
        final CompletableFuture<Result<SqlQueryOk, Err>> f = this.queryClient.sqlQuery(req, Context.newDefault());

        final Result<SqlQueryOk, Err> r = f.get();

        Assert.assertFalse(r.isOk());

        final Err err = r.getErr();

        Assert.assertEquals(ep, err.getErrTo());
        Assert.assertEquals(Collections.singletonList("query_test_table"), err.getFailedTables());
    }

    @Test
    public void queryByArrowFullTypeTest() throws ExecutionException, InterruptedException, IOException {
        final Result<SqlQueryOk, Err> r = queryByArrow();
        Assert.assertTrue(r.isOk());
        final SqlQueryOk queryOk = r.getOk();
        queryOk.stream().forEach(this::checkFullTypeRow);
    }

    private void checkFullTypeRow(final Row row) {
        Assert.assertEquals(14, row.getColumnCount());

        Assert.assertEquals("test", row.getColumn("fString").getValue().getString());
        Assert.assertEquals(Boolean.TRUE, row.getColumn("fBool").getValue().getBoolean());
        Assert.assertEquals(0.64, row.getColumn("fDouble").getValue().getDouble(), 0.000001);
        Assert.assertEquals(0.32f, row.getColumn("fFloat").getValue().getFloat(), 0.000001);
        Assert.assertEquals(-64, row.getColumn("fInt64").getValue().getInt64());
        Assert.assertEquals(-32, row.getColumn("fInt32").getValue().getInt32());
        Assert.assertEquals(-16, row.getColumn("fInt16").getValue().getInt16());
        Assert.assertEquals(-8, row.getColumn("fInt8").getValue().getInt8());
        Assert.assertEquals(64, row.getColumn("fUint64").getValue().getUInt64());
        Assert.assertEquals(32, row.getColumn("fUint32").getValue().getUInt32());
        Assert.assertEquals(16, row.getColumn("fUint16").getValue().getUInt16());
        Assert.assertEquals(8, row.getColumn("fUint8").getValue().getUInt8());
        Assert.assertEquals(1675345488158L, row.getColumn("fTimestamp").getValue().getTimestamp());
        Assert.assertArrayEquals(new byte[] { 1, 2, 3 }, row.getColumn("fVarbinary").getValue().getVarbinary());
    }

    @Test
    public void queryByArrowNullTypeTest() throws ExecutionException, InterruptedException, IOException {
        final Storage.SqlQueryResponse resp = mockSimpleQueryResponse(1, true);
        final Endpoint ep = Endpoint.of("127.0.0.1", 8081);

        Mockito.when(this.routerClient.invoke(Mockito.eq(ep), Mockito.any(), Mockito.any())) //
                .thenReturn(Utils.completedCf(resp));
        Mockito.when(this.routerClient.routeFor(Mockito.any(), Mockito.any())) //
                .thenReturn(Utils.completedCf(new HashMap<>()));
        Mockito.when(this.routerClient.clusterRoute()) //
                .thenReturn(Route.of(ep));

        final SqlQueryRequest req = SqlQueryRequest.newBuilder().forTables("query_test_table") //
                .sql("select * from query_test_table") //
                .build();
        final CompletableFuture<Result<SqlQueryOk, Err>> f = this.queryClient.sqlQuery(req, Context.newDefault());

        final Result<SqlQueryOk, Err> r = f.get();

        Assert.assertTrue(r.isOk());

        final SqlQueryOk queryOk = r.getOk();

        Row row = queryOk.stream().findFirst().get();
        Assert.assertEquals("tvtest", row.getColumn("t1").getValue().getString());
        Assert.assertTrue(row.getColumn("f1").getValue().isNull());
    }

    private Result<SqlQueryOk, Err> queryByArrow() throws IOException, ExecutionException, InterruptedException {
        final Storage.SqlQueryResponse resp = mockAllTypeQueryResponse(1);
        final Endpoint ep = Endpoint.of("127.0.0.1", 8081);

        Mockito.when(this.routerClient.invoke(Mockito.eq(ep), Mockito.any(), Mockito.any())) //
                .thenReturn(Utils.completedCf(resp));
        Mockito.when(this.routerClient.routeFor(Mockito.any(), Mockito.any())) //
                .thenReturn(Utils.completedCf(new HashMap<>()));
        Mockito.when(this.routerClient.clusterRoute()) //
                .thenReturn(Route.of(ep));

        final SqlQueryRequest req = SqlQueryRequest.newBuilder().forTables("query_test_table") //
                .sql("select number from query_test_table") //
                .build();
        final CompletableFuture<Result<SqlQueryOk, Err>> f = this.queryClient.sqlQuery(req, Context.newDefault());

        return f.get();
    }

    private SqlQueryOk mockQueryOk() {
        try {
            Thread.sleep(20);
            return queryByArrow().unwrapOr(SqlQueryOk.emptyOk());
        } catch (final Throwable t) {
            ThrowUtil.throwException(t);
        }
        return null;
    }

    private Storage.SqlQueryResponse mockSimpleQueryResponse(int rowCount, boolean hasNullField) throws IOException {
        final Common.ResponseHeader header = Common.ResponseHeader.newBuilder() //
                .setCode(Result.SUCCESS) //
                .build();

        BufferAllocator allocator = new RootAllocator();
        VarCharVector varCharVector = new VarCharVector("t1", allocator);
        IntVector intVector = new IntVector("f1", allocator);

        for (int i = 0; i < rowCount; i++) {
            varCharVector.setSafe(i, "tvtest".getBytes());
            if (hasNullField) {
                intVector.setNull(i);
            } else {
                intVector.setSafe(i, 123);
            }
        }
        varCharVector.setValueCount(rowCount);
        intVector.setValueCount(rowCount);

        List<Field> fields = Arrays.asList(varCharVector.getField(), intVector.getField());
        List<FieldVector> vectors = Arrays.asList(varCharVector, intVector);
        VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors);

        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ArrowStreamWriter writer = new ArrowStreamWriter(root, /*DictionaryProvider=*/null, out);
        writer.writeBatch();
        writer.close();

        Storage.ArrowPayload arrowPayload = Storage.ArrowPayload.newBuilder()
                .setCompression(Storage.ArrowPayload.Compression.NONE)
                .addRecordBatches(ByteStringHelper.wrap(out.toByteArray())).build();

        return Storage.SqlQueryResponse.newBuilder() //
                .setHeader(header) //
                .setArrow(arrowPayload).build();
    }

    private Storage.SqlQueryResponse mockAllTypeQueryResponse(int rowCount) throws IOException {
        final Common.ResponseHeader header = Common.ResponseHeader.newBuilder() //
                .setCode(Result.SUCCESS) //
                .build();

        BufferAllocator allocator = new RootAllocator();
        VarCharVector stringVector = new VarCharVector("fString", allocator);
        BitVector boolVector = new BitVector("fBool", allocator);
        Float8Vector doubleVector = new Float8Vector("fDouble", allocator);
        Float4Vector floatVector = new Float4Vector("fFloat", allocator);
        BigIntVector int64Vector = new BigIntVector("fInt64", allocator);
        IntVector int32Vector = new IntVector("fInt32", allocator);
        SmallIntVector int16Vector = new SmallIntVector("fInt16", allocator);
        TinyIntVector int8Vector = new TinyIntVector("fInt8", allocator);
        UInt8Vector uint64Vector = new UInt8Vector("fUint64", allocator);
        UInt4Vector uint32Vector = new UInt4Vector("fUint32", allocator);
        UInt2Vector uint16Vector = new UInt2Vector("fUint16", allocator);
        UInt1Vector uint8Vector = new UInt1Vector("fUint8", allocator);
        TimeStampMilliVector tsVector = new TimeStampMilliVector("fTimestamp", allocator);
        VarBinaryVector varBinaryVector = new VarBinaryVector("fVarbinary", allocator);

        for (int i = 0; i < rowCount; i++) {
            stringVector.setSafe(i, "test".getBytes());
            boolVector.setSafe(i, 1);
            doubleVector.setSafe(i, 0.64);
            floatVector.setSafe(i, 0.32f);
            int64Vector.setSafe(i, -64);
            int32Vector.setSafe(i, -32);
            int16Vector.setSafe(i, -16);
            int8Vector.setSafe(i, -8);
            uint64Vector.setSafe(i, 64);
            uint32Vector.setSafe(i, 32);
            uint16Vector.setSafe(i, 16);
            uint8Vector.setSafe(i, 8);
            tsVector.setSafe(i, 1675345488158L);
            varBinaryVector.setSafe(i, new byte[] { 1, 2, 3 });
        }
        stringVector.setValueCount(rowCount);
        boolVector.setValueCount(rowCount);
        doubleVector.setValueCount(rowCount);
        floatVector.setValueCount(rowCount);
        int64Vector.setValueCount(rowCount);
        int32Vector.setValueCount(rowCount);
        int16Vector.setValueCount(rowCount);
        int8Vector.setValueCount(rowCount);
        uint64Vector.setValueCount(rowCount);
        uint32Vector.setValueCount(rowCount);
        uint16Vector.setValueCount(rowCount);
        uint8Vector.setValueCount(rowCount);
        tsVector.setValueCount(rowCount);
        varBinaryVector.setValueCount(rowCount);

        List<Field> fields = Arrays.asList(stringVector.getField(), boolVector.getField(), doubleVector.getField(),
                floatVector.getField(), int64Vector.getField(), int32Vector.getField(), int16Vector.getField(),
                int8Vector.getField(), uint64Vector.getField(), uint32Vector.getField(), uint16Vector.getField(),
                uint8Vector.getField(), tsVector.getField(), varBinaryVector.getField());

        List<FieldVector> vectors = Arrays.asList(stringVector, boolVector, doubleVector, floatVector, int64Vector,
                int32Vector, int16Vector, int8Vector, uint64Vector, uint32Vector, uint16Vector, uint8Vector, tsVector,
                varBinaryVector);
        VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors);

        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ArrowStreamWriter writer = new ArrowStreamWriter(root, /*DictionaryProvider=*/null, out);
        writer.writeBatch();
        writer.close();

        Storage.ArrowPayload arrowPayload = Storage.ArrowPayload.newBuilder()
                .setCompression(Storage.ArrowPayload.Compression.NONE)
                .addRecordBatches(ByteStringHelper.wrap(out.toByteArray())).build();

        return Storage.SqlQueryResponse.newBuilder() //
                .setHeader(header) //
                .setArrow(arrowPayload).build();
    }

    @Test
    public void blockingStreamQueryTest() {
        final SqlQueryOk data = mockQueryOk();
        final int respCount = 10;
        final BlockingStreamIterator streams = new BlockingStreamIterator(1000, TimeUnit.MILLISECONDS);
        final Observer<SqlQueryOk> obs = streams.getObserver();
        new Thread(() -> {
            for (int i = 0; i < respCount; i++) {
                if (i == 0) {
                    obs.onNext(data);
                } else {
                    obs.onNext(mockQueryOk());
                }
            }
            obs.onCompleted();
        }).start();

        final Iterator<Row> it = new RowIterator(streams);

        int i = 0;
        while (it.hasNext()) {
            i++;
            checkFullTypeRow(it.next());
        }
        Assert.assertEquals(respCount, i);
    }

    @Test(expected = IteratorException.class)
    public void blockingStreamQueryOnErrTest() {
        final int respCount = 10;
        final BlockingStreamIterator streams = new BlockingStreamIterator(1000, TimeUnit.MILLISECONDS);
        final Observer<SqlQueryOk> obs = streams.getObserver();
        new Thread(() -> {
            for (int i = 0; i < respCount; i++) {
                if (i == 5) {
                    obs.onError(new StreamException("fail to stream query"));
                } else {
                    obs.onNext(mockQueryOk());
                }
            }
            obs.onCompleted();
        }).start();

        final Iterator<Row> it = new RowIterator(streams);

        int i = 0;
        while (it.hasNext()) {
            i++;
            checkFullTypeRow(it.next());
        }
        Assert.assertEquals(respCount, i);
    }
}