// Licensed to the LF AI & Data foundation 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 <utility>
#include <cstdlib>
#include "folly/futures/Future.h"
#include "storage/FileWriter.h"

namespace milvus::storage {

FileWriter::FileWriter(std::string filename, io::Priority priority)
    : filename_(std::move(filename)),
      priority_(priority),
      rate_limiter_(io::WriteRateLimiter::GetInstance()) {
    // high priority always use buffered mode, otherwise use the global mode
    auto mode =
        priority_ == io::Priority::HIGH ? WriteMode::BUFFERED : GetMode();

    use_direct_io_ = mode == WriteMode::DIRECT;
    use_writer_pool_ = FileWriteWorkerPool::GetInstance().HasPool();

    // allocate an internal aligned buffer for both modes to batch writes
    size_t buf_size = GetBufferSize();
    AssertInfo(
        buf_size != 0 && (buf_size % ALIGNMENT_BYTES) == 0,
        "Buffer size must be greater than 0 and aligned to the alignment "
        "size, buf_size: {}, alignment size: {}",
        buf_size,
        ALIGNMENT_BYTES);
    capacity_ = buf_size;
    aligned_buf_ = aligned_alloc(ALIGNMENT_BYTES, capacity_);
    if (aligned_buf_ == nullptr) {
        ThrowInfo(ErrorCode::MemAllocateFailed,
                  "Failed to allocate aligned buffer of size {}",
                  capacity_);
    }

    auto open_flags = O_CREAT | O_RDWR | O_TRUNC;
    if (use_direct_io_) {
#ifndef __APPLE__
        open_flags |= O_DIRECT;
#endif
    }

    fd_ = open(filename_.c_str(), open_flags, S_IRUSR | S_IWUSR);
    if (fd_ == -1) {
        Cleanup();
        ThrowInfo(ErrorCode::FileCreateFailed,
                  "Failed to open file: {}, error: {}",
                  filename_,
                  strerror(errno));
    }

#ifdef __APPLE__
    if (use_direct_io_) {
        auto ret = fcntl(fd_, F_NOCACHE, 1);
        if (ret == -1) {
            Cleanup();
            ThrowInfo(ErrorCode::FileCreateFailed,
                      "Failed to set F_NOCACHE on file: {}, error: {}",
                      filename_,
                      strerror(errno));
        }
    }
#endif
}

FileWriter::~FileWriter() {
    Cleanup();
}

void
FileWriter::Cleanup() noexcept {
    if (fd_ != -1) {
        close(fd_);
        fd_ = -1;
    }
    if (aligned_buf_ != nullptr) {
        free(aligned_buf_);
        aligned_buf_ = nullptr;
    }
}

bool
FileWriter::PositionedWrite(const void* data,
                            size_t nbyte,
                            size_t file_offset) {
    const char* src = static_cast<const char*>(data);
    size_t left = nbyte;

    while (left != 0) {
        ssize_t done = pwrite(fd_, src, left, file_offset);
        if (done < 0) {
            if (errno == EINTR) {
                continue;
            }
            return false;
        }
        left -= done;
        file_offset += done;
        src += done;
    }

    return true;
}

void
FileWriter::PositionedWriteWithCheck(const void* data,
                                     size_t nbyte,
                                     size_t file_offset) {
    size_t bytes_to_write = nbyte;
    int32_t empty_loops = 0;
    int64_t total_wait_us = 0;
    size_t alignment_bytes = use_direct_io_ ? ALIGNMENT_BYTES : 1;
    while (bytes_to_write != 0) {
        auto allowed_bytes =
            rate_limiter_.Acquire(bytes_to_write, alignment_bytes, priority_);
        if (allowed_bytes == 0) {
            ++empty_loops;
            // if the empty loops is too large or the total wait time is too long, we should write the data directly
            if (empty_loops > MAX_EMPTY_LOOPS || total_wait_us > MAX_WAIT_US) {
                allowed_bytes = bytes_to_write;
                empty_loops = 0;
                total_wait_us = 0;
            } else {
                int64_t wait_us = (1 << (empty_loops / 10)) *
                                  rate_limiter_.GetRateLimitPeriod();
                std::this_thread::sleep_for(std::chrono::microseconds(wait_us));
                total_wait_us += wait_us;
                continue;
            }
        }
        if (!PositionedWrite(data, allowed_bytes, file_offset)) {
            Cleanup();
            ThrowInfo(ErrorCode::FileWriteFailed,
                      "Failed to write to file: {}, error: {}",
                      filename_,
                      strerror(errno));
        }
        file_offset += allowed_bytes;
        bytes_to_write -= allowed_bytes;
        data = static_cast<const char*>(data) + allowed_bytes;
    }
}

void
FileWriter::WriteInternal(const void* data, size_t nbyte) {
    const char* src = static_cast<const char*>(data);
    size_t left_size = nbyte;

    // we should fill and handle the cached aligned buffer first
    if (offset_ != 0) {
        size_t cpy_size = capacity_ - offset_;
        memcpy(static_cast<char*>(aligned_buf_) + offset_, src, cpy_size);
        PositionedWriteWithCheck(aligned_buf_, capacity_, file_size_);
        file_size_ += capacity_;
        left_size -= cpy_size;
        src += cpy_size;
        offset_ = 0;
    }

    // if the left data is aligned, we can just write it to the file and only save the tail to the aligned buffer
    // it will save the time of memcpy
    if (reinterpret_cast<uintptr_t>(src) % ALIGNMENT_BYTES == 0) {
        size_t aligned_left_size = left_size & ~ALIGNMENT_MASK;
        left_size -= aligned_left_size;
        while (aligned_left_size != 0) {
            size_t bytes_written = std::min(aligned_left_size, capacity_);
            PositionedWriteWithCheck(src, bytes_written, file_size_);
            file_size_ += bytes_written;
            aligned_left_size -= bytes_written;
            src += bytes_written;
        }
    }

    // finally, handle the left unaligned data by the aligned buffer
    while (left_size >= capacity_) {
        size_t copy_size = capacity_ - offset_;
        memcpy(static_cast<char*>(aligned_buf_) + offset_, src, copy_size);
        PositionedWriteWithCheck(aligned_buf_, capacity_, file_size_);
        file_size_ += capacity_;
        offset_ = 0;
        left_size -= copy_size;
        src += copy_size;
    }

    // save the tail to the aligned buffer
    if (left_size > 0) {
        memcpy(static_cast<char*>(aligned_buf_) + offset_, src, left_size);
        offset_ += left_size;
        src += left_size;
    }

    assert(src == static_cast<const char*>(data) + nbyte);
}

void
FileWriter::Write(const void* data, size_t nbyte) {
    // if the data can fit in the aligned buffer, we can just copy it to the aligned buffer
    if (nbyte <= capacity_ - offset_) {
        const char* src = static_cast<const char*>(data);
        memcpy(static_cast<char*>(aligned_buf_) + offset_, src, nbyte);
        offset_ += nbyte;
        return;
    }

    if (!use_writer_pool_) {
        WriteInternal(data, nbyte);
        return;
    }

    auto promise = std::make_shared<folly::Promise<folly::Unit>>();
    auto future = promise->getFuture();
    auto task = [this, data, nbyte, promise]() {
        try {
            WriteInternal(data, nbyte);
            promise->setValue(folly::Unit{});
        } catch (...) {
            promise->setException(
                folly::exception_wrapper(std::current_exception()));
        }
    };

    // try to add the task to the writer pool
    // fallback to write the data directly if the task cannot be added to the pool
    if (FileWriteWorkerPool::GetInstance().AddTask(task)) {
        try {
            future.wait();
        } catch (const std::exception& e) {
            Cleanup();
            ThrowInfo(ErrorCode::FileWriteFailed,
                      "Failed to write to file: {}, error: {}",
                      filename_,
                      e.what());
        }
    } else {
        WriteInternal(data, nbyte);
    }
}

void
FileWriter::FlushWithDirectIO() {
    size_t nearest_aligned_offset =
        (offset_ + ALIGNMENT_MASK) & ~ALIGNMENT_MASK;
    memset(static_cast<char*>(aligned_buf_) + offset_,
           0,
           nearest_aligned_offset - offset_);
    PositionedWriteWithCheck(aligned_buf_, nearest_aligned_offset, file_size_);
    file_size_ += offset_;
    // truncate the file to the actual size since the file written by the aligned buffer may be larger than the actual size
    if (ftruncate(fd_, file_size_) != 0) {
        Cleanup();
        ThrowInfo(ErrorCode::FileWriteFailed,
                  "Failed to truncate file: {}, error: {}",
                  filename_,
                  strerror(errno));
    }
    offset_ = 0;
}

void
FileWriter::FlushWithBufferedIO() {
    if (offset_ == 0) {
        return;
    }
    PositionedWriteWithCheck(aligned_buf_, offset_, file_size_);
    file_size_ += offset_;
    offset_ = 0;
}

size_t
FileWriter::Finish() {
    // if the aligned buffer is not empty, we should flush it to the file
    if (offset_ != 0) {
        auto promise = std::make_shared<folly::Promise<folly::Unit>>();
        auto future = promise->getFuture();
        auto task = [this, promise]() {
            try {
                if (use_direct_io_) {
                    FlushWithDirectIO();
                } else {
                    FlushWithBufferedIO();
                }
                promise->setValue(folly::Unit{});
            } catch (...) {
                promise->setException(
                    folly::exception_wrapper(std::current_exception()));
            }
        };

        if (FileWriteWorkerPool::GetInstance().AddTask(task)) {
            try {
                future.wait();
            } catch (const std::exception& e) {
                Cleanup();
                ThrowInfo(ErrorCode::FileWriteFailed,
                          "Failed to flush file: {}, error: {}",
                          filename_,
                          e.what());
            }
        } else {
            if (use_direct_io_) {
                FlushWithDirectIO();
            } else {
                FlushWithBufferedIO();
            }
        }
    }

    // clean up the file writer
    Cleanup();

    // return the file size
    return file_size_;
}

FileWriter::WriteMode FileWriter::mode_ = FileWriter::WriteMode::BUFFERED;
size_t FileWriter::buffer_size_ = DEFAULT_BUFFER_SIZE;

void
FileWriter::SetMode(WriteMode mode) {
    if (mode != WriteMode::BUFFERED && mode != WriteMode::DIRECT) {
        LOG_WARN(
            "Invalid write mode: {}, expected: BUFFERED or DIRECT, "
            "set to BUFFERED",
            static_cast<int>(mode));
        mode = WriteMode::BUFFERED;
    }
    mode_ = mode;
    LOG_INFO("Set write mode to {}", static_cast<uint8_t>(mode));
}

void
FileWriter::SetBufferSize(size_t buffer_size) {
    if (buffer_size > MAX_BUFFER_SIZE) {
        LOG_WARN("Invalid buffer size: {}, expected: <= {}, set to {}",
                 buffer_size,
                 MAX_BUFFER_SIZE,
                 MAX_BUFFER_SIZE);
        buffer_size = MAX_BUFFER_SIZE;
    } else if (buffer_size < MIN_BUFFER_SIZE) {
        LOG_WARN("Invalid buffer size: {}, expected: >= {}, set to {}",
                 buffer_size,
                 MIN_BUFFER_SIZE,
                 MIN_BUFFER_SIZE);
        buffer_size = MIN_BUFFER_SIZE;
    } else {
        buffer_size = (buffer_size + ALIGNMENT_MASK) & ~ALIGNMENT_MASK;
    }
    buffer_size_ = buffer_size;
    LOG_INFO("Set buffer size to {}", buffer_size);
}

FileWriter::WriteMode
FileWriter::GetMode() {
    return mode_;
}

size_t
FileWriter::GetBufferSize() {
    return buffer_size_;
}

}  // namespace milvus::storage
