#pragma once

#include "storage/buffer_manager/spill_result.h"
#include "storage/table/column_chunk_data.h"
#include "storage/table/update_info.h"

namespace kuzu {
namespace storage {
class MemoryManager;

struct ChunkCheckpointState {
    std::unique_ptr<ColumnChunkData> chunkData;
    common::row_idx_t startRow;
    common::length_t numRows;

    ChunkCheckpointState(std::unique_ptr<ColumnChunkData> chunkData, common::row_idx_t startRow,
        common::length_t numRows)
        : chunkData{std::move(chunkData)}, startRow{startRow}, numRows{numRows} {}
};

struct ColumnCheckpointState {
    ColumnChunkData& persistentData;
    std::vector<ChunkCheckpointState> chunkCheckpointStates;
    common::row_idx_t endRowIdxToWrite;

    ColumnCheckpointState(ColumnChunkData& persistentData,
        std::vector<ChunkCheckpointState> chunkCheckpointStates)
        : persistentData{persistentData}, chunkCheckpointStates{std::move(chunkCheckpointStates)},
          endRowIdxToWrite{0} {
        for (const auto& chunkCheckpointState : this->chunkCheckpointStates) {
            const auto endRowIdx = chunkCheckpointState.startRow + chunkCheckpointState.numRows;
            if (endRowIdx > endRowIdxToWrite) {
                endRowIdxToWrite = endRowIdx;
            }
        }
    }
};

class ColumnChunk {
public:
    ColumnChunk(MemoryManager& mm, common::LogicalType&& dataType, uint64_t capacity,
        bool enableCompression, ResidencyState residencyState, bool initializeToZero = true);
    ColumnChunk(MemoryManager& mm, common::LogicalType&& dataType, bool enableCompression,
        ColumnChunkMetadata metadata);
    ColumnChunk(MemoryManager& mm, bool enableCompression, std::unique_ptr<ColumnChunkData> data);

    void initializeScanState(ChunkState& state, const Column* column) const;
    void scan(const transaction::Transaction* transaction, const ChunkState& state,
        common::ValueVector& output, common::offset_t offsetInChunk, common::length_t length) const;
    template<ResidencyState SCAN_RESIDENCY_STATE>
    void scanCommitted(const transaction::Transaction* transaction, ChunkState& chunkState,
        ColumnChunk& output, common::row_idx_t startRow = 0,
        common::row_idx_t numRows = common::INVALID_ROW_IDX) const;
    void lookup(const transaction::Transaction* transaction, const ChunkState& state,
        common::offset_t rowInChunk, common::ValueVector& output,
        common::sel_t posInOutputVector) const;
    void update(const transaction::Transaction* transaction, common::offset_t offsetInChunk,
        const common::ValueVector& values);

    uint64_t getEstimatedMemoryUsage() const {
        return getResidencyState() == ResidencyState::ON_DISK ? 0 : data->getEstimatedMemoryUsage();
    }
    void serialize(common::Serializer& serializer) const;
    static std::unique_ptr<ColumnChunk> deserialize(MemoryManager& mm, common::Deserializer& deSer);

    uint64_t getNumValues() const { return data->getNumValues(); }
    void setNumValues(const uint64_t numValues) const { data->setNumValues(numValues); }

    common::row_idx_t getNumUpdatedRows(const transaction::Transaction* transaction) const;

    std::pair<std::unique_ptr<ColumnChunk>, std::unique_ptr<ColumnChunk>> scanUpdates(
        const transaction::Transaction* transaction) const;

    void setData(std::unique_ptr<ColumnChunkData> data) { this->data = std::move(data); }
    // Functions to access the in memory data.
    ColumnChunkData& getData() const { return *data; }
    const ColumnChunkData& getConstData() const { return *data; }
    std::unique_ptr<ColumnChunkData> moveData() { return std::move(data); }

    common::LogicalType& getDataType() { return data->getDataType(); }
    const common::LogicalType& getDataType() const { return data->getDataType(); }
    bool isCompressionEnabled() const { return enableCompression; }

    ResidencyState getResidencyState() const { return data->getResidencyState(); }
    bool hasUpdates() const { return updateInfo != nullptr; }
    bool hasUpdates(const transaction::Transaction* transaction, common::row_idx_t startRow,
        common::length_t numRows) const;
    // These functions should only work on in-memory and temporary column chunks.
    void resetToEmpty() const { data->resetToEmpty(); }
    void resetToAllNull() const { data->resetToAllNull(); }
    void resize(uint64_t newSize) const { data->resize(newSize); }
    void resetUpdateInfo() {
        if (updateInfo) {
            updateInfo.reset();
        }
    }

    void loadFromDisk() { data->loadFromDisk(); }
    SpillResult spillToDisk() { return data->spillToDisk(); }

    MergedColumnChunkStats getMergedColumnChunkStats(
        const transaction::Transaction* transaction) const;

    void reclaimStorage(PageAllocator& pageAllocator) const;

private:
    void scanCommittedUpdates(const transaction::Transaction* transaction, ColumnChunkData& output,
        common::offset_t startOffsetInOutput, common::row_idx_t startRowScanned,
        common::row_idx_t numRows) const;

private:
    MemoryManager& mm;
    // TODO(Guodong): This field should be removed. Ideally it shouldn't be cached anywhere in
    // storage structures, instead should be fed into functions needed from ClientContext dbConfig.
    bool enableCompression;
    std::unique_ptr<ColumnChunkData> data;
    // Update versions.
    std::unique_ptr<UpdateInfo> updateInfo;
};

} // namespace storage
} // namespace kuzu
