/*
 * 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.
 */

#include "paimon/common/file_index/rangebitmap/range_bitmap.h"

#include <algorithm>

#include "fmt/format.h"
#include "paimon/common/file_index/rangebitmap/dictionary/chunked_dictionary.h"
#include "paimon/common/file_index/rangebitmap/dictionary/key_factory.h"
#include "paimon/common/file_index/rangebitmap/utils/literal_serialization_utils.h"
#include "paimon/common/io/data_output_stream.h"
#include "paimon/common/io/memory_segment_output_stream.h"
#include "paimon/common/memory/memory_segment_utils.h"
#include "paimon/common/utils/field_type_utils.h"
#include "paimon/io/data_input_stream.h"
#include "paimon/memory/bytes.h"

namespace paimon {

Result<std::unique_ptr<RangeBitmap>> RangeBitmap::Create(
    const std::shared_ptr<InputStream>& input_stream, const int64_t offset,
    const FieldType field_type, const std::shared_ptr<MemoryPool>& pool) {
    PAIMON_RETURN_NOT_OK(input_stream->Seek(offset, SeekOrigin::FS_SEEK_SET));
    const auto data_in = std::make_shared<DataInputStream>(input_stream);
    PAIMON_ASSIGN_OR_RAISE(int32_t header_length, data_in->ReadValue<int32_t>());
    PAIMON_ASSIGN_OR_RAISE(int8_t version, data_in->ReadValue<int8_t>());
    if (version != kCurrentVersion) {
        return Status::Invalid(fmt::format("RangeBitmap unsupported version: {}", version));
    }
    PAIMON_ASSIGN_OR_RAISE(int32_t rid, data_in->ReadValue<int32_t>());
    PAIMON_ASSIGN_OR_RAISE(int32_t cardinality, data_in->ReadValue<int32_t>());
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<KeyFactory> shared_key_factory,
                           KeyFactory::Create(field_type));
    PAIMON_ASSIGN_OR_RAISE(LiteralSerDeUtils::Deserializer key_deserializer,
                           LiteralSerDeUtils::CreateValueReader(field_type));
    auto min = Literal{field_type};
    auto max = Literal{field_type};
    if (cardinality > 0) {
        PAIMON_ASSIGN_OR_RAISE(min, key_deserializer(data_in, pool.get()));
        PAIMON_ASSIGN_OR_RAISE(max, key_deserializer(data_in, pool.get()));
    }
    PAIMON_ASSIGN_OR_RAISE(int32_t dictionary_length, data_in->ReadValue<int32_t>());
    auto dictionary_offset = static_cast<int32_t>(offset + sizeof(int32_t) + header_length);
    int32_t bsi_offset = dictionary_offset + dictionary_length;
    return std::unique_ptr<RangeBitmap>(new RangeBitmap(rid, cardinality, dictionary_offset,
                                                        bsi_offset, min, max, shared_key_factory,
                                                        input_stream, pool));
}

Status RangeBitmap::Not(RoaringBitmap32* out) {
    out->Flip(0, rid_);
    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 is_not_null, this->IsNotNull());
    *out &= is_not_null;
    return Status::OK();
}

Result<RoaringBitmap32> RangeBitmap::Eq(const Literal& key) {
    if (cardinality_ <= 0) {
        return RoaringBitmap32();
    }
    PAIMON_ASSIGN_OR_RAISE(int32_t min_compare, key.CompareTo(min_));
    PAIMON_ASSIGN_OR_RAISE(int32_t max_compare, key.CompareTo(max_));
    PAIMON_ASSIGN_OR_RAISE(auto* bit_slice_ptr, this->GetBitSliceIndex());
    if (min_compare == 0 && max_compare == 0) {
        return bit_slice_ptr->IsNotNull({});
    }
    if (min_compare < 0 || max_compare > 0) {
        return RoaringBitmap32();
    }
    PAIMON_ASSIGN_OR_RAISE(auto* dictionary, this->GetDictionary());
    PAIMON_ASSIGN_OR_RAISE(int32_t code, dictionary->Find(key));
    if (code < 0) {
        return RoaringBitmap32();
    }
    return bit_slice_ptr->Eq(code);
}

Result<RoaringBitmap32> RangeBitmap::Neq(const Literal& key) {
    if (cardinality_ <= 0) {
        return RoaringBitmap32();
    }
    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 result, Eq(key));
    PAIMON_RETURN_NOT_OK(Not(&result));
    return result;
}

Result<RoaringBitmap32> RangeBitmap::Lt(const Literal& key) {
    if (cardinality_ <= 0) {
        return RoaringBitmap32();
    }
    PAIMON_ASSIGN_OR_RAISE(int32_t min_compare, key.CompareTo(min_));
    PAIMON_ASSIGN_OR_RAISE(int32_t max_compare, key.CompareTo(max_));
    if (max_compare > 0) {
        return IsNotNull();
    }
    if (min_compare <= 0) {
        return RoaringBitmap32();
    }
    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 result, Gte(key));
    PAIMON_RETURN_NOT_OK(Not(&result));
    return result;
}

Result<RoaringBitmap32> RangeBitmap::Lte(const Literal& key) {
    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 lt_result, Lt(key));
    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 eq_result, Eq(key));
    lt_result |= eq_result;
    return lt_result;
}

Result<RoaringBitmap32> RangeBitmap::Gt(const Literal& key) {
    if (cardinality_ <= 0) {
        return RoaringBitmap32();
    }
    PAIMON_ASSIGN_OR_RAISE(int32_t max_compare, key.CompareTo(max_));
    if (max_compare >= 0) {
        return RoaringBitmap32();
    }
    PAIMON_ASSIGN_OR_RAISE(int32_t min_compare, key.CompareTo(min_));
    if (min_compare < 0) {
        return IsNotNull();
    }
    PAIMON_ASSIGN_OR_RAISE(auto* dictionary, this->GetDictionary());
    PAIMON_ASSIGN_OR_RAISE(int32_t code, dictionary->Find(key));
    PAIMON_ASSIGN_OR_RAISE(auto* bit_slice_ptr, this->GetBitSliceIndex());
    if (code >= 0) {
        return bit_slice_ptr->Gt(code);
    }
    return bit_slice_ptr->Gte(-code - 1);
}

Result<RoaringBitmap32> RangeBitmap::Gte(const Literal& key) {
    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 gt_result, Gt(key));
    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 eq_result, Eq(key));
    gt_result |= eq_result;
    return gt_result;
}

Result<RoaringBitmap32> RangeBitmap::In(const std::vector<Literal>& keys) {
    if (cardinality_ <= 0) {
        return RoaringBitmap32();
    }
    RoaringBitmap32 result{};
    for (const auto& key : keys) {
        PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 bitmap, Eq(key));
        result |= bitmap;
    }
    return result;
}

Result<RoaringBitmap32> RangeBitmap::NotIn(const std::vector<Literal>& keys) {
    if (cardinality_ <= 0) {
        return RoaringBitmap32();
    }
    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 result, In(keys));
    PAIMON_RETURN_NOT_OK(Not(&result));
    return result;
}

Result<RoaringBitmap32> RangeBitmap::IsNull() {
    if (cardinality_ <= 0) {
        if (rid_ > 0) {
            RoaringBitmap32 result;
            result.AddRange(0, rid_);
            return result;
        }
        return RoaringBitmap32();
    }

    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 non_null_bitmap, IsNotNull());
    non_null_bitmap.Flip(0, rid_);
    return non_null_bitmap;
}

Result<RoaringBitmap32> RangeBitmap::IsNotNull() {
    if (cardinality_ <= 0) {
        return RoaringBitmap32();
    }
    PAIMON_ASSIGN_OR_RAISE(auto* bit_slice_ptr, this->GetBitSliceIndex());
    PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 result, bit_slice_ptr->IsNotNull({}));
    return result;
}

RangeBitmap::RangeBitmap(const int32_t rid, const int32_t cardinality,
                         const int32_t dictionary_offset, const int32_t bsi_offset,
                         const Literal& min, const Literal& max,
                         const std::shared_ptr<KeyFactory>& key_factory,
                         const std::shared_ptr<InputStream>& input_stream,
                         const std::shared_ptr<MemoryPool>& pool)
    : pool_(pool),
      rid_(rid),
      cardinality_(cardinality),
      bsi_offset_(bsi_offset),
      dictionary_offset_(dictionary_offset),
      min_(min),
      max_(max),
      key_factory_(key_factory),
      input_stream_(input_stream),
      bsi_(nullptr),
      dictionary_(nullptr) {}

Result<BitSliceIndexBitmap*> RangeBitmap::GetBitSliceIndex() {
    if (bsi_ == nullptr) {
        PAIMON_ASSIGN_OR_RAISE(bsi_,
                               BitSliceIndexBitmap::Create(input_stream_, bsi_offset_, pool_));
    }
    return bsi_.get();
}

Result<Dictionary*> RangeBitmap::GetDictionary() {
    if (dictionary_ == nullptr) {
        PAIMON_ASSIGN_OR_RAISE(
            dictionary_, ChunkedDictionary::Create(key_factory_->GetFieldType(), input_stream_,
                                                   dictionary_offset_, pool_));
    }
    return dictionary_.get();
}

Result<std::unique_ptr<RangeBitmap::Appender>> RangeBitmap::Appender::Create(
    const std::shared_ptr<KeyFactory>& factory, int64_t limited_serialized_size_in_bytes,
    const std::shared_ptr<MemoryPool>& pool) {
    return std::unique_ptr<Appender>(new Appender(factory, limited_serialized_size_in_bytes, pool));
}

RangeBitmap::Appender::Appender(const std::shared_ptr<KeyFactory>& factory,
                                const int64_t limited_serialized_size_in_bytes,
                                const std::shared_ptr<MemoryPool>& pool)
    : pool_(pool),
      rid_(0),
      bitmaps_(LiteralComparator(factory)),
      factory_(factory),
      chunk_size_bytes_limit_(limited_serialized_size_in_bytes) {}

void RangeBitmap::Appender::Append(const Literal& key) {
    if (!key.IsNull()) {
        bitmaps_[key].Add(rid_);
    }
    rid_++;
}

Result<PAIMON_UNIQUE_PTR<Bytes>> RangeBitmap::Appender::Serialize() const {
    int32_t code = 0;
    const int32_t max_code = bitmaps_.empty() ? 0 : static_cast<int32_t>(bitmaps_.size() - 1);
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BitSliceIndexBitmap::Appender> bsi_appender,
                           BitSliceIndexBitmap::Appender::Create(0, max_code, pool_));
    if (chunk_size_bytes_limit_ >= std::numeric_limits<int32_t>::max()) {
        return Status::Invalid(fmt::format(
            "Chunk size cannot be larger than 2GB, current bytes: {}", chunk_size_bytes_limit_));
    }
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ChunkedDictionary::Appender> dictionary_appender,
                           ChunkedDictionary::Appender::Create(
                               factory_, static_cast<int32_t>(chunk_size_bytes_limit_), pool_));
    for (const auto& [key, bitmap] : bitmaps_) {
        PAIMON_RETURN_NOT_OK(dictionary_appender->AppendSorted(key, code));
        for (auto it = bitmap.Begin(); it != bitmap.End(); ++it) {
            PAIMON_RETURN_NOT_OK(bsi_appender->Append(*it, code));
        }
        code++;
    }
    PAIMON_ASSIGN_OR_RAISE(LiteralSerDeUtils::Serializer serializer,
                           LiteralSerDeUtils::CreateValueWriter(factory_->GetFieldType()));
    Literal min{factory_->GetFieldType()};
    Literal max{factory_->GetFieldType()};
    if (!bitmaps_.empty()) {
        min = bitmaps_.begin()->first;
        max = bitmaps_.rbegin()->first;
    }
    PAIMON_ASSIGN_OR_RAISE(int32_t min_size, LiteralSerDeUtils::GetSerializedSizeInBytes(min));
    PAIMON_ASSIGN_OR_RAISE(int32_t max_size, LiteralSerDeUtils::GetSerializedSizeInBytes(max));
    int32_t header_size = 0;
    header_size += sizeof(int8_t);               // version
    header_size += sizeof(int32_t);              // rid
    header_size += sizeof(int32_t);              // cardinality
    header_size += min.IsNull() ? 0 : min_size;  // min literal size
    header_size += max.IsNull() ? 0 : max_size;  // max literal size
    header_size += sizeof(int32_t);              // dictionary length
    PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR<Bytes> dictionary_bytes,
                           dictionary_appender->Serialize());
    auto dictionary_length = static_cast<int32_t>(dictionary_bytes->size());
    PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR<Bytes> bsi_bytes, bsi_appender->Serialize());
    size_t bsi_length = bsi_bytes->size();
    const auto data_output_stream = std::make_shared<MemorySegmentOutputStream>(
        MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_);
    data_output_stream->WriteValue<int32_t>(header_size);
    data_output_stream->WriteValue<int8_t>(kCurrentVersion);
    data_output_stream->WriteValue<int32_t>(rid_);
    data_output_stream->WriteValue<int32_t>(static_cast<int32_t>(bitmaps_.size()));
    if (!min.IsNull()) {
        PAIMON_RETURN_NOT_OK(serializer(data_output_stream, min));
    }
    if (!max.IsNull()) {
        PAIMON_RETURN_NOT_OK(serializer(data_output_stream, max));
    }
    data_output_stream->WriteValue<int32_t>(dictionary_length);
    data_output_stream->Write(dictionary_bytes->data(), dictionary_length);
    data_output_stream->Write(bsi_bytes->data(), bsi_length);
    return MemorySegmentUtils::CopyToBytes(data_output_stream->Segments(), 0,
                                           static_cast<int32_t>(data_output_stream->CurrentSize()),
                                           pool_.get());
}
}  // namespace paimon
