#include "catch.hpp"
#include "test_helpers.hpp"
#include "duckdb/parser/parsed_data/create_table_function_info.hpp"

using namespace duckdb;
using namespace std;

// Dummy TableInOutFunction that:
// - sums all INTEGER values in each row
// - only emits 1 row per call to ThrottlingSum::Function, caching the remainder
// - during flushing of caching operators still emits only 1 row sum per call, meaning that multiple flushes are
// required to correctly process this operator
struct ThrottlingSum {
	struct ThrottlingSumLocalData : public LocalTableFunctionState {
		ThrottlingSumLocalData() {
		}
		duckdb::vector<int> row_sums;
		idx_t current_idx = 0;
	};

	static duckdb::unique_ptr<GlobalTableFunctionState> ThrottlingSumGlobalInit(ClientContext &context,
	                                                                            TableFunctionInitInput &input) {
		return make_uniq<GlobalTableFunctionState>();
	}

	static duckdb::unique_ptr<LocalTableFunctionState> ThrottlingSumLocalInit(ExecutionContext &context,
	                                                                          TableFunctionInitInput &input,
	                                                                          GlobalTableFunctionState *global_state) {
		return make_uniq<ThrottlingSumLocalData>();
	}

	static duckdb::unique_ptr<FunctionData> Bind(ClientContext &context, TableFunctionBindInput &input,
	                                             duckdb::vector<LogicalType> &return_types,
	                                             duckdb::vector<string> &names) {
		return_types.emplace_back(LogicalType::INTEGER);
		names.emplace_back("total");
		return make_uniq<TableFunctionData>();
	}

	static OperatorResultType Function(ExecutionContext &context, TableFunctionInput &data_p, DataChunk &input,
	                                   DataChunk &output) {
		auto &local_state = data_p.local_state->Cast<ThrottlingSum::ThrottlingSumLocalData>();

		for (idx_t row_idx = 0; row_idx < input.size(); row_idx++) {
			int sum = 0;
			for (idx_t col_idx = 0; col_idx < input.ColumnCount(); col_idx++) {
				if (input.data[col_idx].GetType() == LogicalType::INTEGER) {
					sum += input.data[col_idx].GetValue(row_idx).GetValue<int>();
				}
			}
			local_state.row_sums.push_back(sum);
		}

		if (PhysicalOperator::SelectOperatorCachingMode(context) == OperatorCachingMode::UNORDERED) {
			// Caching is allowed
			if (local_state.current_idx < local_state.row_sums.size()) {
				output.SetCardinality(1);
				output.SetValue(0, 0, Value(local_state.row_sums[local_state.current_idx++]));
			} else {
				output.SetCardinality(0);
			}
		} else {
			// Caching is not allowed, we should emit everything!
			auto to_emit = local_state.row_sums.size() - local_state.current_idx;
			for (idx_t i = 0; i < to_emit; i++) {
				output.SetValue(0, i, Value(local_state.row_sums[local_state.current_idx + i]));
			}
			local_state.current_idx += to_emit;
			output.SetCardinality(to_emit);
		}

		return OperatorResultType::NEED_MORE_INPUT;
	}

	static OperatorFinalizeResultType Finalize(ExecutionContext &context, TableFunctionInput &data_p,
	                                           DataChunk &output) {
		auto &local_state = data_p.local_state->Cast<ThrottlingSum::ThrottlingSumLocalData>();

		if (local_state.current_idx < local_state.row_sums.size()) {
			output.SetCardinality(1);
			output.SetValue(0, 0, Value(local_state.row_sums[local_state.current_idx++]));
			return OperatorFinalizeResultType::HAVE_MORE_OUTPUT;
		} else {
			return OperatorFinalizeResultType::FINISHED;
		}
	}

	static void Register(Connection &con) {
		// Create our test TableFunction
		con.BeginTransaction();
		auto &client_context = *con.context;
		auto &catalog = Catalog::GetSystemCatalog(client_context);
		TableFunction caching_table_in_out("throttling_sum", {LogicalType::TABLE}, nullptr, ThrottlingSum::Bind,
		                                   ThrottlingSum::ThrottlingSumGlobalInit,
		                                   ThrottlingSum::ThrottlingSumLocalInit);
		caching_table_in_out.in_out_function = ThrottlingSum::Function;
		caching_table_in_out.in_out_function_final = ThrottlingSum::Finalize;
		CreateTableFunctionInfo caching_table_in_out_info(caching_table_in_out);
		catalog.CreateTableFunction(*con.context, caching_table_in_out_info);
		con.Commit();
	}
};

struct LateralStructEcho {
	static duckdb::unique_ptr<FunctionData> Bind(ClientContext &context, TableFunctionBindInput &input,
	                                             duckdb::vector<LogicalType> &return_types,
	                                             duckdb::vector<string> &names) {
		return_types.emplace_back(LogicalType::BIGINT);
		names.emplace_back("outer_i");
		return_types.emplace_back(LogicalType::BIGINT);
		names.emplace_back("limit_value");
		return_types.emplace_back(LogicalType::VARCHAR);
		names.emplace_back("label_value");
		return make_uniq<TableFunctionData>();
	}

	static OperatorResultType Function(ExecutionContext &context, TableFunctionInput &data_p, DataChunk &input,
	                                   DataChunk &output) {
		for (idx_t row_idx = 0; row_idx < input.size(); row_idx++) {
			auto struct_value = input.data[0].GetValue(row_idx);
			auto &children = StructValue::GetChildren(struct_value);
			output.SetValue(0, row_idx, children[0]);
			output.SetValue(1, row_idx, children[1]);
			output.SetValue(2, row_idx, children[2]);
		}
		output.SetCardinality(input.size());
		return OperatorResultType::NEED_MORE_INPUT;
	}

	static void Register(Connection &con) {
		con.BeginTransaction();
		auto &client_context = *con.context;
		auto &catalog = Catalog::GetSystemCatalog(client_context);
		auto struct_type = LogicalType::STRUCT(
		    {{"outer_i", LogicalType::BIGINT}, {"limit", LogicalType::BIGINT}, {"label", LogicalType::VARCHAR}});
		TableFunction lateral_struct_echo("lateral_struct_echo", {struct_type}, nullptr, LateralStructEcho::Bind);
		lateral_struct_echo.in_out_function = LateralStructEcho::Function;
		CreateTableFunctionInfo lateral_struct_echo_info(lateral_struct_echo);
		catalog.CreateTableFunction(*con.context, lateral_struct_echo_info);
		con.Commit();
	}
};

TEST_CASE("Caching TableInOutFunction", "[filter][.]") {
	DuckDB db(nullptr);
	Connection con(db);

	ThrottlingSum::Register(con);

	// Check result
	auto result2 =
	    con.Query("SELECT * FROM throttling_sum((select i::INTEGER, (i+1)::INTEGER as j from range(0,3) tbl(i)));");
	REQUIRE(result2->ColumnCount() == 1);
	REQUIRE(CHECK_COLUMN(result2, 0, {1, 3, 5}));

	// TODO: streaming these is currently unsupported

	// Large result into aggregation
	auto result3 = con.Query(
	    "SELECT sum(total) FROM throttling_sum((select i::INTEGER, (i+1)::INTEGER as j from range(0,130000) tbl(i)));");
	REQUIRE(result3->ColumnCount() == 1);
	REQUIRE(CHECK_COLUMN(result3, 0, {Value::BIGINT(16900000000)}));
}

TEST_CASE("Parallel execution with caching table in out functions", "[filter][.]") {
	DuckDB db(nullptr);
	Connection con(db);

	ThrottlingSum::Register(con);

	auto result = con.Query("CREATE TABLE test_data as select i::INTEGER from range(0,200000) tbl(i);");
	auto result2 = con.Query("SELECT * FROM throttling_sum((select * from test_data));");

	REQUIRE(result2->ColumnCount() == 1);
	REQUIRE(result2->RowCount() == 200000);
	REQUIRE(CHECK_COLUMN(result2, 0, {0, 1, 2, 3, 4, 5}));
}

TEST_CASE("Lateral table in out function preserves constant struct fields", "[tablefunction]") {
	DuckDB db(nullptr);
	Connection con(db);

	LateralStructEcho::Register(con);

	auto result = con.Query(R"(
		SELECT echoed.outer_i, echoed.limit_value, echoed.label_value
		FROM range(3) outer_rows(i)
		CROSS JOIN LATERAL lateral_struct_echo({'outer_i': i, 'limit': 1, 'label': 'fixed'}) AS echoed
		ORDER BY echoed.outer_i
	)");
	if (result->HasError()) {
		INFO(result->GetError());
	}
	REQUIRE(!result->HasError());
	REQUIRE(result->ColumnCount() == 3);
	REQUIRE(CHECK_COLUMN(result, 0, {0, 1, 2}));
	REQUIRE(CHECK_COLUMN(result, 1, {1, 1, 1}));
	REQUIRE(CHECK_COLUMN(result, 2, {"fixed", "fixed", "fixed"}));
}
