# 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.
#

# for package
import binascii
import logging

import numpy as np
import pandas as pd
from thrift.transport import TTransport
from iotdb.thrift.rpc.IClientRPCService import TSFetchResultsReq, TSCloseOperationReq
from iotdb.utils.IoTDBConstants import TSDataType

logger = logging.getLogger("IoTDB")


class IoTDBRpcDataSet(object):
    TIMESTAMP_STR = "Time"
    # VALUE_IS_NULL = "The value got by %s (column name) is NULL."
    START_INDEX = 2
    FLAG = 0x80

    def __init__(
        self,
        sql,
        column_name_list,
        column_type_list,
        column_name_index,
        ignore_timestamp,
        query_id,
        client,
        statement_id,
        session_id,
        query_data_set,
        fetch_size,
    ):
        self.__statement_id = statement_id
        self.__session_id = session_id
        self.__ignore_timestamp = ignore_timestamp
        self.__sql = sql
        self.__query_id = query_id
        self.__client = client
        self.__fetch_size = fetch_size
        self.__column_size = len(column_name_list)
        self.__default_time_out = 1000

        self.__column_name_list = []
        self.__column_type_list = []
        self.__column_ordinal_dict = {}
        if not ignore_timestamp:
            self.__column_name_list.append(IoTDBRpcDataSet.TIMESTAMP_STR)
            self.__column_type_list.append(TSDataType.INT64)
            self.__column_ordinal_dict[IoTDBRpcDataSet.TIMESTAMP_STR] = 1

        if column_name_index is not None:
            self.__column_type_deduplicated_list = [
                None for _ in range(len(column_name_index))
            ]
            for i in range(len(column_name_list)):
                name = column_name_list[i]
                self.__column_name_list.append(name)
                self.__column_type_list.append(TSDataType[column_type_list[i]])
                if name not in self.__column_ordinal_dict:
                    index = column_name_index[name]
                    self.__column_ordinal_dict[name] = (
                        index + IoTDBRpcDataSet.START_INDEX
                    )
                    self.__column_type_deduplicated_list[index] = TSDataType[
                        column_type_list[i]
                    ]
        else:
            index = IoTDBRpcDataSet.START_INDEX
            self.__column_type_deduplicated_list = []
            for i in range(len(column_name_list)):
                name = column_name_list[i]
                self.__column_name_list.append(name)
                self.__column_type_list.append(TSDataType[column_type_list[i]])
                if name not in self.__column_ordinal_dict:
                    self.__column_ordinal_dict[name] = index
                    index += 1
                    self.__column_type_deduplicated_list.append(
                        TSDataType[column_type_list[i]]
                    )

        self.__time_bytes = bytes(0)
        self.__current_bitmap = [
            bytes(0) for _ in range(len(self.__column_type_deduplicated_list))
        ]
        self.__value = [None for _ in range(len(self.__column_type_deduplicated_list))]
        self.__query_data_set = query_data_set
        self.__is_closed = False
        self.__empty_resultSet = False
        self.__has_cached_record = False
        self.__rows_index = 0

    def close(self):
        if self.__is_closed:
            return
        if self.__client is not None:
            try:
                status = self.__client.closeOperation(
                    TSCloseOperationReq(
                        self.__session_id, self.__query_id, self.__statement_id
                    )
                )
                logger.debug(
                    "close session {}, message: {}".format(
                        self.__session_id, status.message
                    )
                )
            except TTransport.TException as e:
                raise RuntimeError(
                    "close session {} failed because: ".format(self.__session_id), e
                )

            self.__is_closed = True
            self.__client = None

    def next(self):
        if self.has_cached_result():
            self.construct_one_row()
            return True
        if self.__empty_resultSet:
            return False
        if self.fetch_results():
            self.construct_one_row()
            return True
        return False

    def has_cached_result(self):
        return (self.__query_data_set is not None) and (
            len(self.__query_data_set.time) != 0
        )

    def _has_next_result_set(self):
        if self.has_cached_result():
            return True
        if self.__empty_resultSet:
            return False
        if self.fetch_results():
            return True
        return False

    def _to_bitbuffer(self, b):
        return bytes(
            "{:0{}b}".format(int(binascii.hexlify(b), 16), 8 * len(b)), "utf-8"
        )

    def resultset_to_pandas(self):
        result = {}
        for column_name in self.__column_name_list:
            result[column_name] = []
        while self._has_next_result_set():
            time_array = np.frombuffer(
                self.__query_data_set.time, np.dtype(np.longlong).newbyteorder(">")
            )
            if time_array.dtype.byteorder == ">":
                time_array = time_array.byteswap().newbyteorder("<")
            if (
                self.get_ignore_timestamp() is None
                or self.get_ignore_timestamp() is False
            ):
                result[IoTDBRpcDataSet.TIMESTAMP_STR].append(time_array)

            self.__query_data_set.time = []
            total_length = len(time_array)

            for i in range(len(self.__query_data_set.bitmapList)):
                if self.get_ignore_timestamp() is True:
                    column_name = self.get_column_names()[i]
                else:
                    column_name = self.get_column_names()[i + 1]

                location = (
                    self.__column_ordinal_dict[column_name]
                    - IoTDBRpcDataSet.START_INDEX
                )
                if location < 0:
                    continue
                data_type = self.__column_type_deduplicated_list[location]
                value_buffer = self.__query_data_set.valueList[location]
                value_buffer_len = len(value_buffer)

                data_array = None
                if data_type == TSDataType.DOUBLE:
                    data_array = np.frombuffer(
                        value_buffer, np.dtype(np.double).newbyteorder(">")
                    )
                elif data_type == TSDataType.FLOAT:
                    data_array = np.frombuffer(
                        value_buffer, np.dtype(np.float32).newbyteorder(">")
                    )
                elif data_type == TSDataType.BOOLEAN:
                    data_array = np.frombuffer(value_buffer, np.dtype("?"))
                elif data_type == TSDataType.INT32:
                    data_array = np.frombuffer(
                        value_buffer, np.dtype(np.int32).newbyteorder(">")
                    )
                elif data_type == TSDataType.INT64:
                    data_array = np.frombuffer(
                        value_buffer, np.dtype(np.int64).newbyteorder(">")
                    )
                elif data_type == TSDataType.TEXT:
                    j = 0
                    offset = 0
                    data_array = []
                    while offset < value_buffer_len:
                        length = int.from_bytes(
                            value_buffer[offset : offset + 4],
                            byteorder="big",
                            signed=False,
                        )
                        offset += 4
                        value_bytes = value_buffer[offset : offset + length]
                        value = value_bytes.decode("utf-8")
                        data_array.append(value)
                        j += 1
                        offset += length
                    data_array = np.array(data_array, dtype=object)
                else:
                    raise RuntimeError("unsupported data type {}.".format(data_type))
                if data_array.dtype.byteorder == ">":
                    data_array = data_array.byteswap().newbyteorder("<")
                self.__query_data_set.valueList[location] = None

                if len(data_array) < total_length:
                    if data_type == TSDataType.INT32 or data_type == TSDataType.INT64:
                        tmp_array = np.full(total_length, np.nan, np.float32)
                    elif (
                        data_type == TSDataType.FLOAT or data_type == TSDataType.DOUBLE
                    ):
                        tmp_array = np.full(total_length, np.nan, data_array.dtype)
                    elif data_type == TSDataType.BOOLEAN:
                        tmp_array = np.full(total_length, np.nan, np.float32)
                    elif data_type == TSDataType.TEXT:
                        tmp_array = np.full(total_length, None, dtype=data_array.dtype)

                    bitmap_buffer = self.__query_data_set.bitmapList[location]
                    buffer = self._to_bitbuffer(bitmap_buffer)
                    bit_mask = (np.frombuffer(buffer, "u1") - ord("0")).astype(bool)
                    if len(bit_mask) != total_length:
                        bit_mask = bit_mask[:total_length]
                    tmp_array[bit_mask] = data_array

                    if data_type == TSDataType.INT32:
                        tmp_array = pd.Series(tmp_array).astype("Int32")
                    elif data_type == TSDataType.INT64:
                        tmp_array = pd.Series(tmp_array).astype("Int64")
                    elif data_type == TSDataType.BOOLEAN:
                        tmp_array = pd.Series(tmp_array).astype("boolean")

                    data_array = tmp_array

                result[column_name].append(data_array)

        for k, v in result.items():
            if v is None or len(v) < 1 or v[0] is None:
                result[k] = []
            elif v[0].dtype == "Int32":
                result[k] = pd.Series(np.concatenate(v, axis=0)).astype("Int32")
            elif v[0].dtype == "Int64":
                result[k] = pd.Series(np.concatenate(v, axis=0)).astype("Int64")
            elif v[0].dtype == "boolean":
                result[k] = pd.Series(np.concatenate(v, axis=0)).astype("boolean")
            else:
                result[k] = np.concatenate(v, axis=0)

        df = pd.DataFrame(result)
        return df

    def construct_one_row(self):
        # simulating buffer, read 8 bytes from data set and discard first 8 bytes which have been read.
        self.__time_bytes = self.__query_data_set.time[:8]
        self.__query_data_set.time = self.__query_data_set.time[8:]
        for i in range(len(self.__query_data_set.bitmapList)):
            bitmap_buffer = self.__query_data_set.bitmapList[i]

            # another 8 new rows, should move the bitmap buffer position to next byte
            if self.__rows_index % 8 == 0:
                self.__current_bitmap[i] = bitmap_buffer[0]
                self.__query_data_set.bitmapList[i] = bitmap_buffer[1:]
            if not self.is_null(i, self.__rows_index):
                value_buffer = self.__query_data_set.valueList[i]
                data_type = self.__column_type_deduplicated_list[i]

                # simulating buffer
                if data_type == TSDataType.BOOLEAN:
                    self.__value[i] = value_buffer[:1]
                    self.__query_data_set.valueList[i] = value_buffer[1:]
                elif data_type == TSDataType.INT32:
                    self.__value[i] = value_buffer[:4]
                    self.__query_data_set.valueList[i] = value_buffer[4:]
                elif data_type == TSDataType.INT64:
                    self.__value[i] = value_buffer[:8]
                    self.__query_data_set.valueList[i] = value_buffer[8:]
                elif data_type == TSDataType.FLOAT:
                    self.__value[i] = value_buffer[:4]
                    self.__query_data_set.valueList[i] = value_buffer[4:]
                elif data_type == TSDataType.DOUBLE:
                    self.__value[i] = value_buffer[:8]
                    self.__query_data_set.valueList[i] = value_buffer[8:]
                elif data_type == TSDataType.TEXT:
                    length = int.from_bytes(
                        value_buffer[:4], byteorder="big", signed=False
                    )
                    self.__value[i] = value_buffer[4 : 4 + length]
                    self.__query_data_set.valueList[i] = value_buffer[4 + length :]
                else:
                    raise RuntimeError("unsupported data type {}.".format(data_type))
        self.__rows_index += 1
        self.__has_cached_record = True

    def fetch_results(self):
        self.__rows_index = 0
        request = TSFetchResultsReq(
            self.__session_id,
            self.__sql,
            self.__fetch_size,
            self.__query_id,
            True,
            self.__default_time_out,
        )
        try:
            resp = self.__client.fetchResults(request)
            if not resp.hasResultSet:
                self.__empty_resultSet = True
            else:
                self.__query_data_set = resp.queryDataSet
            return resp.hasResultSet
        except TTransport.TException as e:
            raise RuntimeError(
                "Cannot fetch result from server, because of network connection: ", e
            )

    def is_null(self, index, row_num):
        bitmap = self.__current_bitmap[index]
        shift = row_num % 8
        return ((IoTDBRpcDataSet.FLAG >> shift) & (bitmap & 0xFF)) == 0

    def is_null_by_index(self, column_index):
        index = (
            self.__column_ordinal_dict[self.find_column_name_by_index(column_index)]
            - IoTDBRpcDataSet.START_INDEX
        )
        # time column will never be None
        if index < 0:
            return True
        return self.is_null(index, self.__rows_index - 1)

    def is_null_by_name(self, column_name):
        index = self.__column_ordinal_dict[column_name] - IoTDBRpcDataSet.START_INDEX
        # time column will never be None
        if index < 0:
            return True
        return self.is_null(index, self.__rows_index - 1)

    def find_column_name_by_index(self, column_index):
        if column_index <= 0:
            raise Exception("Column index should start from 1")
        if column_index > len(self.__column_name_list):
            raise Exception(
                "column index {} out of range {}".format(
                    column_index, self.__column_size
                )
            )
        return self.__column_name_list[column_index - 1]

    def get_fetch_size(self):
        return self.__fetch_size

    def set_fetch_size(self, fetch_size):
        self.__fetch_size = fetch_size

    def get_column_names(self):
        return self.__column_name_list

    def get_column_types(self):
        return self.__column_type_list

    def get_column_size(self):
        return self.__column_size

    def get_ignore_timestamp(self):
        return self.__ignore_timestamp

    def get_column_ordinal_dict(self):
        return self.__column_ordinal_dict

    def get_column_type_deduplicated_list(self):
        return self.__column_type_deduplicated_list

    def get_values(self):
        return self.__value

    def get_time_bytes(self):
        return self.__time_bytes

    def get_has_cached_record(self):
        return self.__has_cached_record
