// Copyright(C) 2024 InfiniFlow, Inc. All rights reserved.
//
// Licensed 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
//
//     https://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.

module;

module infinity_core:segment_meta.impl;

import :segment_meta;
import :kv_store;
import :kv_code;
import :default_values;
import :table_meta;
import :infinity_exception;
import :block_meta;
import :meta_info;
import :new_catalog;
import :fast_rough_filter;
import :kv_utility;
import :snapshot_info;
import :meta_cache;
import :new_txn;
import :new_txn_manager;

import std;
import third_party;

import column_def;

namespace infinity {

SegmentMeta::SegmentMeta(SegmentID segment_id, TableMeta &table_meta)
    : begin_ts_(table_meta.begin_ts()), commit_ts_(table_meta.commit_ts()), kv_instance_(*table_meta.kv_instance()), table_meta_(table_meta),
      segment_id_(segment_id) {}

Status SegmentMeta::SetFirstDeleteTS(TxnTimeStamp first_delete_ts) {
    std::lock_guard<std::mutex> lock(mtx_);
    std::string first_delete_ts_key = GetSegmentTag("first_delete_ts");
    std::string first_delete_ts_str = fmt::format("{}", first_delete_ts);
    Status status = kv_instance_.Put(first_delete_ts_key, first_delete_ts_str);
    if (!status.ok()) {
        return status;
    }
    first_delete_ts_ = first_delete_ts;
    return Status::OK();
}

Status SegmentMeta::InitSet() {
    {
        Status status = SetFirstDeleteTS(UNCOMMIT_TS);
        if (!status.ok()) {
            return status;
        }
    }
    return Status::OK();
}

// called when restore segment from snapshot with segment meta
// restore segment meta from snapshot with segment meta
Status SegmentMeta::RestoreSet() {
    {
        Status status = SetFirstDeleteTS(first_delete_ts_.value_or(UNCOMMIT_TS));
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

// Status SegmentMeta::UninitSet(UsageFlag usage_flag) { return UninitSet(usage_flag, begin_ts_); }

Status SegmentMeta::UninitSet(UsageFlag usage_flag, TxnTimeStamp begin_ts) {
    {
        std::string block_id_prefix = KeyEncode::CatalogTableSegmentBlockKeyPrefix(table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_);
        auto iter = kv_instance_.GetIterator();
        iter->Seek(block_id_prefix);
        std::vector<std::string> delete_keys;
        while (iter->Valid() && iter->Key().starts_with(block_id_prefix)) {
            TxnTimeStamp commit_ts = std::stoull(iter->Value().ToString());
            if (commit_ts > begin_ts and commit_ts != std::numeric_limits<TxnTimeStamp>::max()) {
                BlockID block_id = std::stoull(iter->Key().ToString().substr(block_id_prefix.size()));
                UnrecoverableError(fmt::format("Block id: {}.{} is not allowed to be removed. commit_ts: {}, begin_ts: {}",
                                               segment_id_,
                                               block_id,
                                               commit_ts,
                                               begin_ts));
            }
            // the key is committed before the txn or the key isn't committed
            delete_keys.push_back(iter->Key().ToString());
            iter->Next();
        }
        for (const std::string &key : delete_keys) {
            Status status = kv_instance_.Delete(key);
            if (!status.ok()) {
                return status;
            }
        }
        block_ids1_.reset();
    }
    {
        std::string first_delete_ts_key = GetSegmentTag("first_delete_ts");
        Status status = kv_instance_.Delete(first_delete_ts_key);
        if (!status.ok()) {
            return status;
        }
    }
    {
        std::string filter_key = GetSegmentTag("fast_rough_filter");
        Status status = kv_instance_.Delete(filter_key);
        if (!status.ok()) {
            if (status.code() != ErrorCode::kNotFound) {
                return status;
            }
        }
    }
    return Status::OK();
}

Status SegmentMeta::LoadFirstDeleteTS() {
    if (first_delete_ts_.has_value()) {
        return Status::OK();
    }

    std::string first_delete_ts_key = GetSegmentTag("first_delete_ts");

    std::shared_ptr<MetaTableCache> table_cache{};
    u64 db_id = table_meta_.db_id();
    const std::string &table_name = table_meta_.table_name();
    MetaCache *meta_cache = table_meta_.meta_cache();
    table_cache = meta_cache->GetTable(db_id, table_name, begin_ts_);
    if (table_cache.get() != nullptr) {
        auto tag_value = table_cache->get_segment_tag(segment_id_, first_delete_ts_key);
        if (tag_value.has_value()) {
            first_delete_ts_ = tag_value;
            return Status::OK();
        }
    }

    std::string first_delete_ts_str;
    Status status = kv_instance_.Get(first_delete_ts_key, first_delete_ts_str);
    if (!status.ok()) {
        return status;
    }
    first_delete_ts_ = std::stoull(first_delete_ts_str);

    if (table_cache.get() != nullptr && table_meta_.txn() != nullptr) {
        table_meta_.txn()->AddCacheInfo(
            std::make_shared<TableCacheSegmentTagInfo>(db_id, table_name, begin_ts_, segment_id_, first_delete_ts_key, *first_delete_ts_));
    }
    return Status::OK();
}

std::string SegmentMeta::GetSegmentTag(const std::string &tag) const {
    return KeyEncode::CatalogTableSegmentTagKey(table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_, tag);
}

TxnTimeStamp SegmentMeta::GetCreateTimestampFromKV() const {
    std::string segment_key = KeyEncode::CatalogTableSegmentKey(table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_);
    std::string create_ts_str;
    Status status = kv_instance_.Get(segment_key, create_ts_str);
    if (!status.ok()) {
        return 0;
    }
    return std::stoull(create_ts_str);
}

Status SegmentMeta::Init() { return Status::OK(); }

Status SegmentMeta::AddBlockWithID(TxnTimeStamp commit_ts, BlockID block_id) {
    Status status;
    std::string block_id_key = KeyEncode::CatalogTableSegmentBlockKey(table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_, block_id);
    std::string commit_ts_str = fmt::format("{}", commit_ts);
    status = kv_instance_.Put(block_id_key, commit_ts_str);
    if (!status.ok()) {
        return status;
    }
    return Status::OK();
}

std::pair<BlockID, Status> SegmentMeta::AddBlockID1(TxnTimeStamp commit_ts) {
    BlockID block_id = 0;
    {
        // Load, take the next id and cache it in one critical section: concurrent appends of the
        // same segment must not hand out the same id, nor reallocate block_ids1_ under a reader.
        std::lock_guard<std::mutex> lock(mtx_);
        if (!block_ids1_) {
            block_ids1_ = infinity::GetTableSegmentBlocks(&kv_instance_,
                                                          table_meta_.db_id_str(),
                                                          table_meta_.table_id_str(),
                                                          segment_id_,
                                                          begin_ts_,
                                                          commit_ts);
        }
        block_id = block_ids1_->empty() ? 0 : block_ids1_->back() + 1;
        block_ids1_->push_back(block_id);
    }

    std::string block_id_key = KeyEncode::CatalogTableSegmentBlockKey(table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_, block_id);
    std::string commit_ts_str = fmt::format("{}", commit_ts);
    Status status = kv_instance_.Put(block_id_key, commit_ts_str);
    if (!status.ok()) {
        return {0, status};
    }
    return {block_id, Status::OK()};
}

Status SegmentMeta::CommitBlock(BlockID block_id, TxnTimeStamp commit_ts) {
    std::string block_id_key = KeyEncode::CatalogTableSegmentBlockKey(table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_, block_id);
    std::string commit_ts_str = fmt::format("{}", commit_ts);
    Status status = kv_instance_.Put(block_id_key, commit_ts_str);
    if (!status.ok()) {
        return status;
    }
    return Status::OK();
}

std::tuple<std::vector<BlockID> *, Status> SegmentMeta::GetBlockIDs1() {
    std::lock_guard<std::mutex> lock(mtx_);
    if (!block_ids1_) {
        block_ids1_ =
            infinity::GetTableSegmentBlocks(&kv_instance_, table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_, begin_ts_, commit_ts_);
    }
    return {&*block_ids1_, Status::OK()};
}

std::vector<BlockID> SegmentMeta::GetBlockIDsCopy() {
    std::lock_guard<std::mutex> lock(mtx_);
    if (!block_ids1_) {
        block_ids1_ =
            infinity::GetTableSegmentBlocks(&kv_instance_, table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_, begin_ts_, commit_ts_);
    }
    // Copy while the lock is held: AddBlockID1 appends to block_ids1_ and may reallocate it.
    return *block_ids1_;
}

std::tuple<size_t, Status> SegmentMeta::GetRowCnt1() {
    std::lock_guard<std::mutex> lock(mtx_);
    if (row_cnt_) {
        return {*row_cnt_, Status::OK()};
    }

    Status status;
#if 1
    row_cnt_ = infinity::GetSegmentRowCount(&kv_instance_, table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_, begin_ts_, commit_ts_);

    return {*row_cnt_, Status::OK()};
#else
    size_t row_cnt = 0;
    {
        std::vector<BlockID> *block_ids_ptr = nullptr;
        std::tie(block_ids_ptr, status) = GetBlockIDs1();
        if (block_ids_ptr->size() > 0) {
            row_cnt += DEFAULT_BLOCK_CAPACITY * (block_ids_ptr->size() - 1);

            BlockMeta block_meta(block_ids_ptr->back(), *this);
            size_t block_row_cnt;
            std::tie(block_row_cnt, status) = block_meta.GetRowCnt1();
            if (!status.ok()) {
                return {0, status};
            }
            row_cnt += block_row_cnt;
        }
    }
    row_cnt_ = row_cnt;
    return {row_cnt, Status::OK()};
#endif
}

Status SegmentMeta::GetFirstDeleteTS(TxnTimeStamp &first_delete_ts) {
    std::lock_guard<std::mutex> lock(mtx_);
    if (!first_delete_ts_) {
        Status status = LoadFirstDeleteTS();
        if (!status.ok()) {
            return status;
        }
    }
    // if (*first_delete_ts_ > begin_ts_) {
    //     first_delete_ts = UNCOMMIT_TS;
    // }
    first_delete_ts = *first_delete_ts_;
    return Status::OK();
}

Status SegmentMeta::CommitFirstDeleteTS(const TxnTimeStamp &commit_ts) {
    std::lock_guard<std::mutex> lock(mtx_);
    if (!first_delete_ts_) {
        Status status = LoadFirstDeleteTS();
        if (!status.ok()) {
            return status;
        }
    }
    if (*first_delete_ts_ == UNCOMMIT_TS && commit_ts != UNCOMMIT_TS) {
        std::string first_delete_ts_key = GetSegmentTag("first_delete_ts");
        std::string first_delete_ts_str = fmt::format("{}", commit_ts);
        Status status = kv_instance_.Put(first_delete_ts_key, first_delete_ts_str);
        if (!status.ok()) {
            return status;
        }
        first_delete_ts_ = commit_ts;
    }
    return Status::OK();
}

std::tuple<std::shared_ptr<SegmentInfo>, Status> SegmentMeta::GetSegmentInfo() {
    auto segment_info = std::make_shared<SegmentInfo>();
    std::shared_ptr<std::vector<std::shared_ptr<ColumnDef>>> column_defs = nullptr;
    i64 column_count = 0;
    SegmentID unsealed_segment_id = 0;
    Status status;
    std::tie(column_defs, status) = table_meta_.GetColumnDefs();
    if (!status.ok()) {
        return {nullptr, status};
    }
    column_count = column_defs->size();
    status = table_meta_.GetUnsealedSegmentID(unsealed_segment_id);
    if (!status.ok()) {
        unsealed_segment_id = INVALID_SEGMENT_ID;
    }

    segment_info->segment_id_ = segment_id_;
    segment_info->status_ =
        (segment_id_ == unsealed_segment_id) ? SegmentStatus::kUnsealed : SegmentStatus::kSealed; // TODO: How to determine other status?
    segment_info->column_count_ = column_count;
    segment_info->row_capacity_ = segment_capacity();
    segment_info->storage_size_ = 0; // TODO: How to determine storage size?
    std::vector<BlockID> *block_ids_ptr = nullptr;
    std::tie(block_ids_ptr, status) = GetBlockIDs1();
    if (!status.ok()) {
        return {nullptr, status};
    }
    segment_info->block_count_ = block_ids_ptr->size();
    std::tie(segment_info->row_count_, status) = GetRowCnt1();
    if (!status.ok()) {
        return {nullptr, status};
    }
    segment_info->actual_row_count_ = segment_info->row_count_;
    for (BlockID block_id : *block_ids_ptr) {
        BlockMeta block_meta(block_id, *this);
        std::vector<std::string> file_paths;
        status = NewCatalog::GetBlockFilePaths(block_meta, file_paths, nullptr);
        if (!status.ok()) {
            return {nullptr, status};
        }
        segment_info->files_.insert(segment_info->files_.end(), file_paths.begin(), file_paths.end());
    }
    return {std::move(segment_info), Status::OK()};
}

Status SegmentMeta::GetFastRoughFilter(std::shared_ptr<FastRoughFilter> &fast_rough_filter) {
    fast_rough_filter.reset();

    std::unique_lock lock(mtx_);

    if (fast_rough_filter_) {
        fast_rough_filter = fast_rough_filter_;
        return Status::OK();
    }

    std::string filter_key = GetSegmentTag("fast_rough_filter");
    std::string filter_str;
    Status status = kv_instance_.Get(filter_key, filter_str);
    if (!status.ok()) {
        return status;
    }
    fast_rough_filter_ = std::make_shared<FastRoughFilter>();
    fast_rough_filter_->DeserializeFromString(filter_str);
    fast_rough_filter = fast_rough_filter_;

    return Status::OK();
}

Status SegmentMeta::SetFastRoughFilter(std::shared_ptr<FastRoughFilter> fast_rough_filter) {
    std::string filter_key = GetSegmentTag("fast_rough_filter");
    std::string filter_str = fast_rough_filter->SerializeToString();
    std::lock_guard<std::mutex> lock(mtx_);
    Status status = kv_instance_.Put(filter_key, filter_str);
    if (!status.ok()) {
        return status;
    }
    fast_rough_filter_ = fast_rough_filter;
    return Status::OK();
}

std::tuple<std::shared_ptr<SegmentSnapshotInfo>, Status> SegmentMeta::MapMetaToSnapShotInfo() {

    std::shared_ptr<SegmentSnapshotInfo> segment_snapshot_info = std::make_shared<SegmentSnapshotInfo>();
    Status status;
    segment_snapshot_info->segment_id_ = segment_id_;
    segment_snapshot_info->create_ts_ = GetCreateTimestampFromKV();
    status = GetFirstDeleteTS(segment_snapshot_info->first_delete_ts_);
    if (!status.ok()) {
        return {nullptr, status};
    }

    std::tie(segment_snapshot_info->row_count_, status) = GetRowCnt1();
    if (!status.ok()) {
        return {nullptr, status};
    }

    std::vector<BlockID> *block_ids_ptr = nullptr;
    std::tie(block_ids_ptr, status) = GetBlockIDs1();
    if (!status.ok()) {
        return {nullptr, status};
    }

    for (BlockID block_id : *block_ids_ptr) {
        BlockMeta block_meta(block_id, *this);
        auto [block_snapshot, block_status] = block_meta.MapMetaToSnapShotInfo();
        if (!block_status.ok()) {
            return {nullptr, block_status};
        }
        segment_snapshot_info->block_snapshots_.push_back(block_snapshot);
    }

    return {std::move(segment_snapshot_info), Status::OK()};
}

Status SegmentMeta::RestoreFromSnapshot(const WalSegmentInfoV2 &segment_info, bool is_link_files) {
    if (!is_link_files) {
        Status status = RestoreSet();
        if (!status.ok()) {
            return status;
        }
    }
    for (const BlockID &block_id : segment_info.block_ids_) {
        if (!is_link_files) {
            Status status = AddBlockWithID(commit_ts(), block_id);
            if (!status.ok()) {
                return status;
            }
        }
        BlockMeta block_meta(block_id, *this);
        Status status = block_meta.RestoreFromSnapshot();
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

} // namespace infinity