/*
 * 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/core/mergetree/levels.h"

#include <algorithm>
namespace paimon {
bool Levels::Level0Comparator::operator()(const std::shared_ptr<DataFileMeta>& a,
                                          const std::shared_ptr<DataFileMeta>& b) const {
    if (a->max_sequence_number != b->max_sequence_number) {
        // file with larger sequence number should be in front
        return a->max_sequence_number > b->max_sequence_number;
    } else {
        // When two or more jobs are writing the same merge tree, it is
        // possible that multiple files have the same maxSequenceNumber. In
        // this case we have to compare their file names so that files with
        // same maxSequenceNumber won't be "de-duplicated" by the tree set.
        int64_t min_seq_a = a->min_sequence_number;
        int64_t min_seq_b = b->min_sequence_number;
        if (min_seq_a != min_seq_b) {
            return min_seq_a < min_seq_b;
        }
        // If minSequenceNumber is also the same, use creation time
        Timestamp time_a = a->creation_time;
        Timestamp time_b = b->creation_time;
        if (time_a != time_b) {
            return time_a < time_b;
        }
        // Final fallback: filename (to ensure uniqueness in set)
        return a->file_name < b->file_name;
    }
}

Result<std::unique_ptr<Levels>> Levels::Create(
    const std::shared_ptr<FieldsComparator>& key_comparator,
    const std::vector<std::shared_ptr<DataFileMeta>>& input_files, int32_t num_levels) {
    // in case the num of levels is not specified explicitly
    int32_t restored_num_levels = -1;
    for (const auto& file : input_files) {
        if (file->level > restored_num_levels) {
            restored_num_levels = file->level;
        }
    }
    restored_num_levels = std::max(restored_num_levels + 1, num_levels);
    if (restored_num_levels <= 1) {
        return Status::Invalid("Number of levels must be at least 2.");
    }

    std::set<std::shared_ptr<DataFileMeta>, Levels::Level0Comparator> level0;
    std::vector<SortedRun> levels;
    levels.reserve(restored_num_levels - 1);
    for (int32_t i = 1; i < restored_num_levels; ++i) {
        levels.push_back(SortedRun::Empty());
    }
    auto level_map = GroupByLevel(input_files);
    for (auto& [level, files] : level_map) {
        PAIMON_RETURN_NOT_OK(
            UpdateLevel(level, /*before=*/{}, /*after=*/files, key_comparator, &levels, &level0));
    }

    size_t total_file_num = level0.size();
    for (const auto& run : levels) {
        total_file_num += run.Files().size();
    }
    if (total_file_num != input_files.size()) {
        return Status::Invalid(
            "Number of files stored in Levels does not equal to the size of inputFiles. This "
            "is unexpected.");
    }
    return std::unique_ptr<Levels>(new Levels(key_comparator, level0, levels));
}

int32_t Levels::NumberOfSortedRuns() const {
    int32_t number_of_runs = level0_.size();
    for (const auto& run : levels_) {
        if (!run.IsEmpty()) {
            number_of_runs++;
        }
    }
    return number_of_runs;
}

void Levels::AddDropFileCallback(DropFileCallback* callback) {
    drop_file_callbacks_.push_back(callback);
}

void Levels::RemoveDropFileCallback(DropFileCallback* callback) {
    drop_file_callbacks_.erase(
        std::remove(drop_file_callbacks_.begin(), drop_file_callbacks_.end(), callback),
        drop_file_callbacks_.end());
}

Status Levels::AddLevel0File(const std::shared_ptr<DataFileMeta>& file) {
    if (file->level != 0) {
        return Status::Invalid("must add level0 file in AddLevel0File");
    }
    level0_.insert(file);
    return Status::OK();
}

int32_t Levels::NonEmptyHighestLevel() const {
    for (int32_t i = levels_.size() - 1; i >= 0; i--) {
        if (!levels_[i].IsEmpty()) {
            return i + 1;
        }
    }
    return level0_.empty() ? -1 : 0;
}

int64_t Levels::TotalFileSize() const {
    int64_t total_size = 0;
    for (const auto& file : level0_) {
        total_size += file->file_size;
    }
    for (const auto& run : levels_) {
        total_size += run.TotalSize();
    }
    return total_size;
}

std::vector<std::shared_ptr<DataFileMeta>> Levels::AllFiles() const {
    std::vector<std::shared_ptr<DataFileMeta>> all_files;
    auto runs = LevelSortedRuns();
    for (const auto& run : runs) {
        all_files.insert(all_files.end(), run.run.Files().begin(), run.run.Files().end());
    }
    return all_files;
}

std::vector<LevelSortedRun> Levels::LevelSortedRuns() const {
    std::vector<LevelSortedRun> runs;
    for (const auto& file : level0_) {
        runs.emplace_back(/*level=*/0, SortedRun::FromSingle(file));
    }
    for (int32_t i = 0; i < static_cast<int32_t>(levels_.size()); i++) {
        const auto& run = levels_[i];
        if (!run.IsEmpty()) {
            runs.emplace_back(/*level=*/i + 1, run);
        }
    }
    return runs;
}

Status Levels::Update(const std::vector<std::shared_ptr<DataFileMeta>>& before,
                      const std::vector<std::shared_ptr<DataFileMeta>>& after) {
    auto grouped_before = GroupByLevel(before);
    auto grouped_after = GroupByLevel(after);
    int32_t number_of_levels = NumberOfLevels();
    for (int32_t i = 0; i < number_of_levels; i++) {
        PAIMON_RETURN_NOT_OK(UpdateLevel(i, grouped_before[i], grouped_after[i], key_comparator_,
                                         &levels_, &level0_));
    }
    if (!drop_file_callbacks_.empty()) {
        std::set<std::string> dropped_files;
        for (const auto& file : before) {
            dropped_files.insert(file->file_name);
        }
        // exclude upgrade files
        for (const auto& file : after) {
            dropped_files.erase(file->file_name);
        }
        for (auto* callback : drop_file_callbacks_) {
            if (!callback) {
                continue;
            }
            for (const auto& file_name : dropped_files) {
                callback->NotifyDropFile(file_name);
            }
        }
    }
    return Status::OK();
}

Status Levels::UpdateLevel(int32_t level, const std::vector<std::shared_ptr<DataFileMeta>>& before,
                           const std::vector<std::shared_ptr<DataFileMeta>>& after,
                           const std::shared_ptr<FieldsComparator>& key_comparator,
                           std::vector<SortedRun>* levels,
                           std::set<std::shared_ptr<DataFileMeta>, Level0Comparator>* level0) {
    if (before.empty() && after.empty()) {
        return Status::OK();
    }
    if (level == 0) {
        for (const auto& file : before) {
            level0->erase(file);
        }
        for (const auto& file : after) {
            level0->insert(file);
        }
    } else {
        PAIMON_ASSIGN_OR_RAISE(SortedRun run, RunOfLevel(level, *levels));
        std::vector<std::shared_ptr<DataFileMeta>> files = run.Files();
        for (const auto& before_file : before) {
            auto iter = std::find_if(files.begin(), files.end(), [&before_file](const auto& cur) {
                return before_file->file_name == cur->file_name;
            });
            if (iter != files.end()) {
                files.erase(iter);
            }
        }
        files.insert(files.end(), after.begin(), after.end());
        PAIMON_ASSIGN_OR_RAISE((*levels)[level - 1],
                               SortedRun::FromUnsorted(files, key_comparator));
    }
    return Status::OK();
}

Result<SortedRun> Levels::RunOfLevel(int32_t level, const std::vector<SortedRun>& levels) {
    if (level <= 0) {
        return Status::Invalid("Level0 does not have one single sorted run.");
    }
    return levels[level - 1];
}

std::map<int32_t, std::vector<std::shared_ptr<DataFileMeta>>> Levels::GroupByLevel(
    const std::vector<std::shared_ptr<DataFileMeta>>& files) {
    std::map<int32_t, std::vector<std::shared_ptr<DataFileMeta>>> level_map;
    for (const auto& file : files) {
        level_map[file->level].push_back(file);
    }
    return level_map;
}

}  // namespace paimon
