//===----------------------------------------------------------------------===//
//                         DuckDB
//
// duckdb/storage/table/row_group_collection.hpp
//
//
//===----------------------------------------------------------------------===//

#pragma once

#include "duckdb/storage/table/row_group.hpp"
#include "duckdb/storage/table/segment_tree.hpp"
#include "duckdb/storage/statistics/column_statistics.hpp"
#include "duckdb/storage/table/table_statistics.hpp"
#include "duckdb/storage/storage_index.hpp"
#include "duckdb/common/enums/index_removal_type.hpp"
#include "duckdb/common/enums/row_group_append_mode.hpp"

namespace duckdb {

struct ParallelTableScanState;
struct ParallelCollectionScanState;
class CreateIndexScanState;
class CollectionScanState;
class PersistentTableData;
class TableDataWriter;
class TableIndexList;
class TableStatistics;
struct TableAppendState;
class DuckTransaction;
class BoundConstraint;
class RowGroupSegmentTree;
class StorageCommitState;
struct ColumnSegmentInfo;
class MetadataManager;
struct VacuumState;
struct CollectionCheckpointState;
struct PersistentCollectionData;
class CheckpointTask;
class TableIOManager;
class DataTable;
class RowGroupIterationHelper;
class TableScanState;

class RowGroupCollection {
public:
	RowGroupCollection(shared_ptr<DataTableInfo> info, TableIOManager &io_manager, vector<LogicalType> types,
	                   idx_t row_start, idx_t total_rows = 0);
	RowGroupCollection(shared_ptr<DataTableInfo> info, BlockManager &block_manager, vector<LogicalType> types,
	                   idx_t row_start, idx_t total_rows, idx_t row_group_size);

public:
	idx_t GetTotalRows() const;
	idx_t GetRowGroupCount() const;
	Allocator &GetAllocator() const;

	void Initialize(PersistentCollectionData &data);
	void Initialize(PersistentTableData &data);
	void InitializeEmpty();
	void FinalizeCheckpoint(MetaBlockPointer pointer, const vector<MetaBlockPointer> &existing_pointers);

	bool IsEmpty() const;

	void AppendRowGroup(SegmentLock &l, idx_t start_row);
	//! Get the nth row-group, negative numbers start from the back (so -1 is the last row group, etc)
	optional_ptr<RowGroup> GetRowGroup(int64_t index);
	//! Overrides a row group - should only be used if you know what you're doing (will likely be removed in the future)
	void SetRowGroup(int64_t index, shared_ptr<RowGroup> new_row_group);
	void Verify();
	void Destroy();

	void InitializeScan(const QueryContext &context, CollectionScanState &state, const vector<StorageIndex> &column_ids,
	                    optional_ptr<TableFilterSet> table_filters);
	void InitializeCreateIndexScan(CreateIndexScanState &state);
	void InitializeScanWithOffset(const QueryContext &context, CollectionScanState &state,
	                              const vector<StorageIndex> &column_ids, idx_t start_row, idx_t end_row);
	static bool InitializeScanInRowGroup(ClientContext &context, CollectionScanState &state,
	                                     RowGroupCollection &collection, SegmentNode<RowGroup> &row_group,
	                                     idx_t vector_index, idx_t max_row);
	void InitializeParallelScan(ParallelCollectionScanState &state);
	bool NextParallelScan(ClientContext &context, ParallelCollectionScanState &state, CollectionScanState &scan_state);

	RowGroupIterationHelper Chunks(DuckTransaction &transaction);
	RowGroupIterationHelper Chunks(DuckTransaction &transaction, const vector<StorageIndex> &column_ids);

	void Fetch(TransactionData transaction, DataChunk &result, const vector<StorageIndex> &column_ids,
	           const Vector &row_identifiers, idx_t fetch_count, ColumnFetchState &state);

	//! Returns true, if the row group can fetch the row id for the transaction.
	bool CanFetch(TransactionData, const row_t row_id);

	//! Initialize an append of a variable number of rows. FinalizeAppend must be called after appending is done.
	void InitializeAppend(TableAppendState &state);
	//! Initialize an append with a variable number of rows. FinalizeAppend should not be called after appending is
	//! done.
	void InitializeAppend(TransactionData transaction, TableAppendState &state);
	//! Appends to the row group collection. Returns true if a new row group has been created to append to
	bool Append(DataChunk &chunk, TableAppendState &state);
	//! FinalizeAppend flushes an append with a variable number of rows.
	void FinalizeAppend(TransactionData transaction, TableAppendState &state);
	void CommitAppend(transaction_t commit_id, idx_t row_start, idx_t count);
	void RevertAppendInternal(idx_t start_row);
	void CleanupAppend(transaction_t lowest_transaction, idx_t start, idx_t count);

	void MergeStorage(RowGroupCollection &data, optional_ptr<DataTable> table,
	                  optional_ptr<StorageCommitState> commit_state);
	bool IsPersistent() const;

	void RemoveFromIndexes(const QueryContext &context, TableIndexList &indexes, Vector &row_identifiers, idx_t count,
	                       IndexRemovalType removal_type, optional_idx active_checkpoint = optional_idx());

	idx_t Delete(TransactionData transaction, DataTable &table, row_t *ids, idx_t count);
	void Update(TransactionData transaction, DataTable &table, row_t *ids, const vector<PhysicalIndex> &column_ids,
	            DataChunk &updates);
	void UpdateColumn(TransactionData transaction, DataTable &table, Vector &row_ids,
	                  const vector<column_t> &column_path, DataChunk &updates);

	void Checkpoint(TableDataWriter &writer, TableStatistics &global_stats);

	void InitializeVacuumState(CollectionCheckpointState &checkpoint_state, VacuumState &state,
	                           optional_idx checkpoint_row_group_count);
	bool ScheduleVacuumTasks(CollectionCheckpointState &checkpoint_state, VacuumState &state, idx_t segment_idx,
	                         bool schedule_vacuum);
	unique_ptr<CheckpointTask> GetCheckpointTask(CollectionCheckpointState &checkpoint_state, idx_t segment_idx);

	void CommitDropColumn(const idx_t column_index);
	void CommitDropTable();

	vector<PartitionStatistics> GetPartitionStats() const;
	vector<ColumnSegmentInfo> GetColumnSegmentInfo(const QueryContext &context) const;
	bool SupportsPerColumnWrites();
	const vector<LogicalType> &GetTypes() const;

	shared_ptr<RowGroupCollection> AddColumn(ClientContext &context, ColumnDefinition &new_column,
	                                         ExpressionExecutor &default_executor);
	shared_ptr<RowGroupCollection> RemoveColumn(idx_t col_idx);
	shared_ptr<RowGroupCollection> AlterType(ClientContext &context, idx_t changed_idx, const LogicalType &target_type,
	                                         vector<StorageIndex> bound_columns, Expression &cast_expr);
	void VerifyNewConstraint(const QueryContext &context, DataTable &parent, const BoundConstraint &constraint);

	void SetStats(TableStatistics &new_stats);
	void CopyStats(TableStatistics &stats);
	unique_ptr<BaseStatistics> CopyStats(const StorageIndex &column_id);
	unique_ptr<BlockingSample> GetSample();
	void SetDistinct(column_t column_id, unique_ptr<DistinctStatistics> distinct_stats);

	AttachedDatabase &GetAttached();
	BlockManager &GetBlockManager() {
		return block_manager;
	}
	MetadataManager &GetMetadataManager();
	DataTableInfo &GetTableInfo() {
		return *info;
	}

	idx_t GetAllocationSize() const {
		return allocation_size;
	}

	idx_t GetRowGroupSize() const {
		return row_group_size;
	}
	void SetRowGroupAppendMode(RowGroupAppendMode mode);
	//! Returns the total amount of segments - use sparingly, as this forces all segments to be loaded
	idx_t GetSegmentCount();

	//! Get a ptr to the raw segment tree. This can be useful for some extensions to have directly exposed.
	shared_ptr<RowGroupSegmentTree> GetRowGroups() const;

private:
	optional_ptr<SegmentNode<RowGroup>> NextUpdateRowGroup(RowGroupSegmentTree &row_groups, row_t *ids, idx_t &pos,
	                                                       idx_t count) const;

	void SetRowGroups(shared_ptr<RowGroupSegmentTree> row_groups);

private:
	//! BlockManager
	BlockManager &block_manager;
	//! The row group size of the row group collection
	const idx_t row_group_size;
	//! The number of rows in the table
	atomic<idx_t> total_rows;
	//! The data table info
	shared_ptr<DataTableInfo> info;
	//! The column types of the row group collection
	vector<LogicalType> types;
	//! Lock held when accessing or modifying the owned_row_groups pointer
	mutable mutex row_group_pointer_lock;
	//! The owning pointer of the segment tree
	shared_ptr<RowGroupSegmentTree> owned_row_groups;
	//! Table statistics
	TableStatistics stats;
	//! Allocation size, only tracked for appends
	atomic<idx_t> allocation_size;
	//! Root metadata pointer, if the collection is loaded from disk
	MetaBlockPointer metadata_pointer;
	//! Other metadata pointers
	vector<MetaBlockPointer> metadata_pointers;
	//! Controls whether the next append creates a new row group or reuses the existing one
	RowGroupAppendMode row_group_append_mode;
};

class RowGroupIterationHelper {
public:
	RowGroupIterationHelper(RowGroupCollection &collection, DuckTransaction &transaction,
	                        vector<StorageIndex> column_ids);

private:
	RowGroupCollection &collection;
	DuckTransaction &transaction;
	vector<StorageIndex> column_ids;

private:
	class RowGroupIterator {
	public:
		RowGroupIterator(optional_ptr<RowGroupCollection> collection, optional_ptr<DuckTransaction> transaction,
		                 const vector<StorageIndex> &column_ids);
		~RowGroupIterator();
		//! enable move constructor
		RowGroupIterator(RowGroupIterator &&other) noexcept;

		optional_ptr<RowGroupCollection> collection;
		optional_ptr<DuckTransaction> transaction;
		unique_ptr<DataChunk> chunk;
		unique_ptr<TableScanState> state;

	public:
		RowGroupIterator &operator++();
		bool operator!=(const RowGroupIterator &other) const;
		DataChunk &operator*() const;
	};

public:
	RowGroupIterator begin(); // NOLINT: match stl API
	RowGroupIterator end();   // NOLINT: match stl API
};

} // namespace duckdb
