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

#pragma once

#include <cstdint>
#include <map>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <string>
#include <unordered_map>
#include <vector>

#include "storage/IndexData.h"
#include "storage/FileManager.h"
#include "storage/LocalChunkManager.h"
#include "common/Consts.h"
#include "storage/Types.h"
#include "storage/ThreadPools.h"

namespace milvus::storage {

class DiskFileManagerImpl : public FileManagerImpl {
 private:
    struct LocalDirState;

 public:
    class LocalDirWriteLease {
     public:
        LocalDirWriteLease() = default;
        LocalDirWriteLease(LocalDirWriteLease&& other) noexcept = default;
        LocalDirWriteLease&
        operator=(LocalDirWriteLease&& other) noexcept;
        LocalDirWriteLease(const LocalDirWriteLease&) = delete;
        LocalDirWriteLease&
        operator=(const LocalDirWriteLease&) = delete;
        ~LocalDirWriteLease();

        explicit operator bool() const {
            return state_ != nullptr;
        }

     private:
        friend class DiskFileManagerImpl;

        explicit LocalDirWriteLease(std::shared_ptr<LocalDirState> state);

        void
        Release() noexcept;

        std::shared_ptr<LocalDirState> state_;
    };

    explicit DiskFileManagerImpl(const FileManagerContext& fileManagerContext);

    ~DiskFileManagerImpl() override;

    bool
    LoadFile(const std::string& filename) noexcept override;

    bool
    AddFile(const std::string& filename) noexcept override;

    bool
    AddFileMeta(const FileMeta& file_meta) override;

    std::optional<bool>
    IsExisted(const std::string& filename) noexcept override;

    bool
    RemoveFile(const std::string& filename) noexcept override;

 public:
    bool
    AddTextLog(const std::string& filename) noexcept;

    bool
    AddJsonSharedIndexLog(const std::string& filename) noexcept;

    bool
    AddJsonStatsMetaLog(const std::string& filename) noexcept;

 public:
    std::string
    GetName() const override {
        return "DiskFileManagerImpl";
    }

    std::string
    GetIndexIdentifier();

    std::string
    GetLocalIndexObjectPrefix();

    std::string
    GetLocalTempIndexObjectPrefix();

    // Similar to GetTextIndexIdentifier, segment_id and field_id is also required.
    std::string
    GetLocalTextIndexPrefix();

    std::string
    GetLocalTempTextIndexPrefix();

    std::string
    GetLocalJsonStatsPrefix();

    std::string
    GetLocalTempJsonStatsPrefix();

    std::string
    GetLocalJsonStatsShreddingPrefix();

    std::string
    GetLocalJsonStatsSharedIndexPrefix();

    std::string
    GetLocalJsonStatsShreddingPath(const std::string& file_name);

    // Used for upload index to remote storage, using this index prefix dir as remote storage directory
    std::string
    GetLocalNgramIndexPrefix();

    // Used for loading index, using this index prefix dir to store index.
    std::string
    GetLocalTempNgramIndexPrefix();

    std::string
    GetRemoteJsonStatsLogPrefix();

    std::string
    GetRemoteJsonStatsShreddingPrefix();

    std::string
    GetRemoteJsonStatsSharedIndexPath(const std::string& file_name,
                                      int64_t slice_num);

    std::string
    GetRemoteJsonStatsMetaPath(const std::string& file_name);

    std::string
    GetLocalJsonStatsMetaPrefix();

    std::string
    GetLocalRawDataObjectPrefix();

    std::map<std::string, int64_t>
    GetRemotePathsToFileSize() const {
        return remote_paths_to_size_;
    }

    std::vector<std::string>
    GetLocalFilePaths() const {
        return local_paths_;
    }

    void
    CacheIndexToDisk(const std::vector<std::string>& remote_files,
                     milvus::proto::common::LoadPriority priority);

    void
    CacheTextLogToDisk(const std::vector<std::string>& remote_files,
                       milvus::proto::common::LoadPriority priority);

    void
    CacheNgramIndexToDisk(const std::vector<std::string>& remote_files,
                          milvus::proto::common::LoadPriority priority);

    void
    RemoveIndexFiles();

    void
    RemoveTextLogFiles();

    void
    RemoveJsonStatsSharedIndexFiles();

    void
    RemoveJsonStatsFiles();

    void
    CacheJsonStatsSharedIndexToDisk(
        const std::vector<std::string>& remote_files,
        milvus::proto::common::LoadPriority priority);

    // Cache meta file to local disk
    // Returns local file path
    // remote_file must be an absolute remote path (basePath + relative file)
    std::string
    CacheJsonStatsMetaToDisk(const std::string& remote_file,
                             milvus::proto::common::LoadPriority priority);

    void
    RemoveNgramIndexFiles();

    void
    RemoveRawDataFiles();

    LocalDirWriteLease
    AcquireLocalDirWriteLease(const std::string& dir);

    void
    AddBatchIndexFiles(const std::string& local_file_name,
                       const std::vector<int64_t>& local_file_offsets,
                       const std::vector<std::string>& remote_files,
                       const std::vector<int64_t>& remote_file_sizes);

    template <typename DataType>
    std::string
    CacheRawDataToDisk(const Config& config);

    std::string
    CacheOptFieldToDisk(const Config& config);

    std::string
    GetRemoteIndexPrefix() const {
        return GetRemoteIndexObjectPrefix();
    }

    size_t
    GetAddedTotalFileSize() const {
        return added_total_file_size_;
    }

    std::string
    GetFileName(const std::string& localfile);

 private:
    int64_t
    GetIndexBuildId() {
        return index_meta_.build_id;
    }

    std::string
    GetRemoteIndexPath(const std::string& file_name, int64_t slice_num) const;

    /**
     * @brief Get the Remote Index Path V2
     * @param file_name; v2 will not split the file with slice_num
     * @return std::string
     */
    std::string
    GetRemoteIndexPathV2(const std::string& file_name) const;

    std::string
    GetRemoteTextLogPath(const std::string& file_name, int64_t slice_num) const;

    std::string
    AppendLocalPathGeneration(const std::string& prefix) const;

    static void
    RemoveLocalDirBestEffort(const std::string& dir) noexcept;

    void
    CloseAndRemoveLocalDir(const std::string& dir) noexcept;

    std::shared_ptr<LocalDirState>
    GetOrCreateLocalDirState(const std::string& dir);

    bool
    AddFileInternal(const std::string& file_name,
                    const std::function<std::string(const std::string&, int)>&
                        get_remote_path) noexcept;

    void
    CacheIndexToDiskInternal(const std::vector<std::string>& remote_files,
                             const std::string& local_index_prefix,
                             milvus::proto::common::LoadPriority priority =
                                 milvus::proto::common::LoadPriority::HIGH);

    template <typename DataType>
    std::string
    cache_raw_data_to_disk_internal(const Config& config);

    template <typename T>
    std::string
    cache_raw_data_to_disk_storage_v2(const Config& config);

    std::string
    cache_opt_field_to_disk_v2(const Config& config);

    std::string
    cache_opt_field_to_disk_v3(const Config& config);

    template <typename DataType>
    void
    cache_raw_data_to_disk_common(
        const FieldDataPtr& field_data,
        const std::shared_ptr<LocalChunkManager>& local_chunk_manager,
        std::string& local_data_path,
        bool& file_created,
        uint32_t& dim,
        int64_t& write_offset,
        std::vector<size_t>* offsets = nullptr);

    inline void
    set_bit(std::vector<uint8_t>& bitmap, int64_t bit_pos) {
        bitmap[bit_pos >> 3] |= (1 << (bit_pos & 0x07));
    }

    void
    write_valid_data_file(
        const std::shared_ptr<LocalChunkManager>& local_chunk_manager,
        const std::string& valid_data_path,
        std::vector<uint8_t>& valid_bitmap,
        uint64_t total_num_rows);

 private:
    // local file path (abs path)
    std::vector<std::string> local_paths_;

    // remote file path
    std::map<std::string, int64_t> remote_paths_to_size_;

    mutable std::mutex local_dir_states_mutex_;
    std::unordered_map<std::string, std::shared_ptr<LocalDirState>>
        local_dir_states_;

    size_t added_total_file_size_ = 0;

    uint64_t file_path_generation_;
};

using DiskANNFileManagerImplPtr = std::shared_ptr<DiskFileManagerImpl>;

}  // namespace milvus::storage
