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

#pragma once

#include <cassert>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "arrow/c/abi.h"
#include "arrow/c/helpers.h"
#include "fmt/format.h"
#include "paimon/common/global_index/global_index_utils.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/file_index/file_index_writer.h"
#include "paimon/global_index/global_index_writer.h"
#include "paimon/global_index/io/global_index_file_writer.h"

namespace paimon {
/// A `GlobalIndexWriter` wrapper for `FileIndexWriter`.
class FileIndexWriterWrapper : public GlobalIndexWriter {
 public:
    FileIndexWriterWrapper(const std::string& index_type,
                           const std::shared_ptr<GlobalIndexFileWriter>& file_manager,
                           const std::shared_ptr<FileIndexWriter>& writer)
        : index_type_(index_type), file_manager_(file_manager), writer_(writer) {}

    Status AddBatch(::ArrowArray* c_arrow_array, std::vector<int64_t>&& relative_row_ids) override {
        PAIMON_RETURN_NOT_OK(
            GlobalIndexUtils::CheckRelativeRowIds(c_arrow_array, relative_row_ids, count_));
        auto length = c_arrow_array->length;
        PAIMON_RETURN_NOT_OK(writer_->AddBatch(c_arrow_array));
        count_ += length;
        return Status::OK();
    }

    Result<std::vector<GlobalIndexIOMeta>> Finish() override {
        if (count_ == 0) {
            return std::vector<GlobalIndexIOMeta>();
        }
        PAIMON_ASSIGN_OR_RAISE(std::string file_name, file_manager_->NewFileName(index_type_));
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<OutputStream> out,
                               file_manager_->NewOutputStream(file_name));
        PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR<Bytes> bytes, writer_->SerializedBytes());

        PAIMON_ASSIGN_OR_RAISE(int64_t actual_size, out->Write(bytes->data(), bytes->size()));
        if (actual_size < 0 || static_cast<size_t>(actual_size) != bytes->size()) {
            return Status::IOError(fmt::format("expect write len {} mismatch actual write len {}",
                                               bytes->size(), actual_size));
        }
        PAIMON_RETURN_NOT_OK(out->Flush());
        PAIMON_RETURN_NOT_OK(out->Close());
        GlobalIndexIOMeta meta(file_manager_->ToPath(file_name), /*file_size=*/bytes->size(),
                               /*metadata=*/nullptr);
        return std::vector<GlobalIndexIOMeta>({meta});
    }

 private:
    std::string index_type_;
    int64_t count_ = 0;
    std::shared_ptr<GlobalIndexFileWriter> file_manager_;
    std::shared_ptr<FileIndexWriter> writer_;
};
}  // namespace paimon
