/*
 * 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.iotdb.jdbc;

import org.apache.iotdb.rpc.TSStatusCode;
import org.apache.iotdb.service.rpc.thrift.*;
import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.Statement;
import java.sql.Types;
import java.time.ZoneId;
import java.util.ArrayList;
import java.util.List;

import static org.mockito.Matchers.any;
import static org.mockito.Mockito.*;

/*
    This class is designed to test the function of TsfileQueryResultSet.
    This class also sheds light on the complete execution process of a query sql from the jdbc perspective.

    The test utilizes the mockito framework to mock responses from an IoTDB server.
    The status of the IoTDB server mocked here is determined by the following sql commands:

    "SET STORAGE GROUP TO root.vehicle",
    "CREATE TIMESERIES root.vehicle.d0.s0 WITH DATATYPE=INT32, ENCODING=RLE",
    "CREATE TIMESERIES root.vehicle.d0.s1 WITH DATATYPE=INT64, ENCODING=RLE",
    "CREATE TIMESERIES root.vehicle.d0.s2 WITH DATATYPE=FLOAT, ENCODING=RLE",
    "insert into root.vehicle.d0(timestamp,s0) values(1,101)",
    "insert into root.vehicle.d0(timestamp,s0) values(2,198)",
    "insert into root.vehicle.d0(timestamp,s0) values(100,99)",
    "insert into root.vehicle.d0(timestamp,s0) values(101,99)",
    "insert into root.vehicle.d0(timestamp,s0) values(102,80)",
    "insert into root.vehicle.d0(timestamp,s0) values(103,99)",
    "insert into root.vehicle.d0(timestamp,s0) values(104,90)",
    "insert into root.vehicle.d0(timestamp,s0) values(105,99)",
    "insert into root.vehicle.d0(timestamp,s0) values(106,99)",
    "insert into root.vehicle.d0(timestamp,s0) values(2,10000)",
    "insert into root.vehicle.d0(timestamp,s0) values(50,10000)",
    "insert into root.vehicle.d0(timestamp,s0) values(1000,22222)",
    "DELETE FROM root.vehicle.d0.s0 WHERE time < 104",
    "UPDATE root.vehicle.d0 SET s0 = 33333 WHERE time < 106 and time > 103",
    "insert into root.vehicle.d0(timestamp,s1) values(1,1101)",
    "insert into root.vehicle.d0(timestamp,s1) values(2,198)",
    "insert into root.vehicle.d0(timestamp,s1) values(100,199)",
    "insert into root.vehicle.d0(timestamp,s1) values(101,199)",
    "insert into root.vehicle.d0(timestamp,s1) values(102,180)",
    "insert into root.vehicle.d0(timestamp,s1) values(103,199)",
    "insert into root.vehicle.d0(timestamp,s1) values(104,190)",
    "insert into root.vehicle.d0(timestamp,s1) values(105,199)",
    "insert into root.vehicle.d0(timestamp,s1) values(2,40000)",
    "insert into root.vehicle.d0(timestamp,s1) values(50,50000)",
    "insert into root.vehicle.d0(timestamp,s1) values(1000,55555)",
    "insert into root.vehicle.d0(timestamp,s2) values(1000,55555)",
    "insert into root.vehicle.d0(timestamp,s2) values(2,2.22)",
    "insert into root.vehicle.d0(timestamp,s2) values(3,3.33)",
    "insert into root.vehicle.d0(timestamp,s2) values(4,4.44)",
    "insert into root.vehicle.d0(timestamp,s2) values(102,10.00)",
    "insert into root.vehicle.d0(timestamp,s2) values(105,11.11)",
    "insert into root.vehicle.d0(timestamp,s2) values(1000,1000.11)",
    "insert into root.vehicle.d0(timestamp,s1) values(2000-01-01T08:00:00+08:00, 100)",
 */

public class IoTDBQueryResultSetTest {

  @Mock
  TSExecuteStatementResp execResp;
  @Mock
  TSOperationHandle operationHandle;
  @Mock
  TSHandleIdentifier handleIdentifier;
  @Mock
  private IoTDBConnection connection;
  @Mock
  private TSIService.Iface client;
  @Mock
  private TS_SessionHandle sessHandle;
  @Mock
  private Statement statement;
  @Mock
  private TSFetchMetadataResp fetchMetadataResp;
  @Mock
  private TSFetchResultsResp fetchResultsResp;

  private TSStatusType successStatus = new TSStatusType(TSStatusCode.SUCCESS_STATUS.getStatusCode(),
      "");
  private TSStatus Status_SUCCESS = new TSStatus(successStatus);
  private ZoneId zoneID = ZoneId.systemDefault();

  @Before
  public void before() throws Exception {
    MockitoAnnotations.initMocks(this);

    statement = new IoTDBStatement(connection, client, sessHandle, zoneID);

    when(connection.isClosed()).thenReturn(false);
    when(client.executeStatement(any(TSExecuteStatementReq.class))).thenReturn(execResp);
    operationHandle.hasResultSet = true;
    operationHandle.operationId = handleIdentifier;
    handleIdentifier.queryId = 1L;
    when(execResp.getOperationHandle()).thenReturn(operationHandle);
    when(operationHandle.getOperationId()).thenReturn(handleIdentifier);
    when(handleIdentifier.getQueryId()).thenReturn(1L);
    when(execResp.getStatus()).thenReturn(Status_SUCCESS);

    when(client.fetchMetadata(any(TSFetchMetadataReq.class))).thenReturn(fetchMetadataResp);
    when(fetchMetadataResp.getStatus()).thenReturn(Status_SUCCESS);

    when(client.fetchResults(any(TSFetchResultsReq.class))).thenReturn(fetchResultsResp);
    when(fetchResultsResp.getStatus()).thenReturn(Status_SUCCESS);

    TSStatus closeResp = new TSStatus();
    closeResp = Status_SUCCESS;
    when(client.closeOperation(any(TSCloseOperationReq.class))).thenReturn(closeResp);
  }

  @SuppressWarnings("resource")
  @Test
  public void testQuery() throws Exception {

    String testSql = "select *,s1,s0,s2 from root.vehicle.d0 where s1 > 190 or s2 < 10.0 "
        + "limit 20 offset 1 slimit 4 soffset 2";

    /*
     * step 1: execute statement
     */
    List<String> columns = new ArrayList<>();
    columns.add("root.vehicle.d0.s2");
    columns.add("root.vehicle.d0.s1");
    columns.add("root.vehicle.d0.s0");
    columns.add("root.vehicle.d0.s2");

    List<String> dataTypeList = new ArrayList<>();
//    //BOOLEAN, INT32, INT64, FLOAT, DOUBLE, TEXT
//    dataTypeList.add(TSDataType.INT64.toString());
    dataTypeList.add("FLOAT");
    dataTypeList.add("INT64");
    dataTypeList.add("INT32");
    dataTypeList.add("FLOAT");

    when(execResp.getColumns()).thenReturn(columns);
    when(execResp.getDataTypeList()).thenReturn(dataTypeList);
    when(execResp.getOperationType()).thenReturn("QUERY");
    doReturn("FLOAT").doReturn("INT64").doReturn("INT32").doReturn("FLOAT").when(fetchMetadataResp)
        .getDataType();

    boolean hasResultSet = statement.execute(testSql);
    Assert.assertTrue(hasResultSet);

    verify(fetchMetadataResp, times(0)).getDataType();

    /*
     * step 2: fetch result
     */
    fetchResultsResp.hasResultSet = true; // at the first time to fetch
    TSQueryDataSet tsQueryDataSet = FakedFirstFetchResult();
    when(fetchResultsResp.getQueryDataSet()).thenReturn(tsQueryDataSet);

    try (ResultSet resultSet = statement.getResultSet()) {
      // check columnInfoMap
      Assert.assertEquals(resultSet.findColumn("Time"), 1);
      Assert.assertEquals(resultSet.findColumn("root.vehicle.d0.s2"), 2);
      Assert.assertEquals(resultSet.findColumn("root.vehicle.d0.s1"), 3);
      Assert.assertEquals(resultSet.findColumn("root.vehicle.d0.s0"), 4);

      ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
      // check columnInfoList
      Assert.assertEquals(resultSetMetaData.getColumnName(1), "Time");
      Assert.assertEquals(resultSetMetaData.getColumnName(2), "root.vehicle.d0.s2");
      Assert.assertEquals(resultSetMetaData.getColumnName(3), "root.vehicle.d0.s1");
      Assert.assertEquals(resultSetMetaData.getColumnName(4), "root.vehicle.d0.s0");
      Assert.assertEquals(resultSetMetaData.getColumnName(5), "root.vehicle.d0.s2");
      // check columnTypeList
      Assert.assertEquals(resultSetMetaData.getColumnType(1), Types.TIMESTAMP);
      Assert.assertEquals(resultSetMetaData.getColumnType(2), Types.FLOAT);
      Assert.assertEquals(resultSetMetaData.getColumnType(3), Types.BIGINT);
      Assert.assertEquals(resultSetMetaData.getColumnType(4), Types.INTEGER);
      Assert.assertEquals(resultSetMetaData.getColumnType(5), Types.FLOAT);
      // check fetched result
      int colCount = resultSetMetaData.getColumnCount();
      StringBuilder resultStr = new StringBuilder();
      for (int i = 1; i < colCount + 1; i++) { // meta title
        resultStr.append(resultSetMetaData.getColumnName(i)).append(",");
      }
      resultStr.append("\n");
      while (resultSet.next()) { // data
        for (int i = 1; i <= colCount; i++) {
          resultStr.append(resultSet.getString(i)).append(",");
        }
        resultStr.append("\n");

        fetchResultsResp.hasResultSet = false; // at the second time to fetch
      }
      String standard =
          "Time,root.vehicle.d0.s2,root.vehicle.d0.s1,root.vehicle.d0.s0,root.vehicle.d0.s2,\n"
              + "2,2.22,40000,null,2.22,\n"
              + "3,3.33,null,null,3.33,\n"
              + "4,4.44,null,null,4.44,\n"
              + "50,null,50000,null,null,\n"
              + "100,null,199,null,null,\n"
              + "101,null,199,null,null,\n"
              + "103,null,199,null,null,\n"
              + "105,11.11,199,33333,11.11,\n"
              + "1000,1000.11,55555,22222,1000.11,\n"; // Note the LIMIT&OFFSET clause takes effect
      Assert.assertEquals(resultStr.toString(), standard);
    }
  }

  // fake the first-time fetched result of 'testSql' from an IoTDB server
  private TSQueryDataSet FakedFirstFetchResult() throws IOException {
    List<TSDataType> tsDataTypeList = new ArrayList<>();
    tsDataTypeList.add(TSDataType.FLOAT); // root.vehicle.d0.s2
    tsDataTypeList.add(TSDataType.INT64); // root.vehicle.d0.s1
    tsDataTypeList.add(TSDataType.INT32); // root.vehicle.d0.s0

    Object[][] input = {
        {1L, null, 1101L, null,},
        {2L, 2.22F, 40000L, null,},
        {3L, 3.33F, null, null,},
        {4L, 4.44F, null, null,},
        {50L, null, 50000L, null,},
        {100L, null, 199L, null,},
        {101L, null, 199L, null,},
        {103L, null, 199L, null,},
        {105L, 11.11F, 199L, 33333,},
        {1000L, 1000.11F, 55555L, 22222,}};

    TSQueryDataSet tsQueryDataSet = new TSQueryDataSet();
    int rowCount = input.length;
    tsQueryDataSet.setRowCount(rowCount);

    int columnNum = tsDataTypeList.size();
    int columnNumWithTime = columnNum + 1;
    DataOutputStream[] dataOutputStreams = new DataOutputStream[columnNumWithTime];
    ByteArrayOutputStream[] byteArrayOutputStreams = new ByteArrayOutputStream[columnNumWithTime];
    for (int i = 0; i < columnNumWithTime; i++) {
      byteArrayOutputStreams[i] = new ByteArrayOutputStream();
      dataOutputStreams[i] = new DataOutputStream(byteArrayOutputStreams[i]);
    }

    int valueOccupation = 0;
    for (int i = 0; i < rowCount; i++) {
      Object[] row = input[i];
      // use columnOutput to write byte array
      dataOutputStreams[0].writeLong((long) row[0]);
      for (int k = 0; k < columnNum; k++) {
        DataOutputStream dataOutputStream = dataOutputStreams[k + 1]; // DO NOT FORGET +1
        Object value = row[1 + k];
        if (value == null) {
          dataOutputStream.writeBoolean(true); // is_empty true
        } else {
          dataOutputStream.writeBoolean(false); // is_empty false
          if (k == 0) { // TSDataType.FLOAT
            dataOutputStream.writeFloat((float) value);
            valueOccupation += 4;
          } else if (k == 1) { // TSDataType.INT64
            dataOutputStream.writeLong((long) value);
            valueOccupation += 8;
          } else { // TSDataType.INT32
            dataOutputStream.writeInt((int) value);
            valueOccupation += 4;
          }
        }
      }
    }

    // calculate total valueOccupation
    valueOccupation += rowCount * 8; // note the timestamp column needn't the boolean is_empty
    valueOccupation += rowCount * columnNum; // for all is_empty

    ByteBuffer valueBuffer = ByteBuffer.allocate(valueOccupation);
    for (ByteArrayOutputStream byteArrayOutputStream : byteArrayOutputStreams) {
      valueBuffer.put(byteArrayOutputStream.toByteArray());
    }
    valueBuffer.flip(); // PAY ATTENTION TO HERE
    tsQueryDataSet.setValues(valueBuffer);
    tsQueryDataSet.setRowCount(rowCount);
    return tsQueryDataSet;
  }
}