#include "duckdb/function/scalar/nested_functions.hpp"
#include "duckdb/function/scalar/struct_functions.hpp"
#include "duckdb/common/case_insensitive_map.hpp"
#include "duckdb/planner/expression_binder.hpp"
#include "duckdb/planner/expression/bound_function_expression.hpp"
#include "duckdb/storage/statistics/struct_stats.hpp"

namespace duckdb {

static void StructConcatFunction(DataChunk &args, ExpressionState &state, Vector &result) {
	auto &result_cols = StructVector::GetEntries(result);
	idx_t offset = 0;

	if (!args.AllConstant()) {
		// Unless all arguments are constant, we flatten the input to make sure it's homogeneous
		args.Flatten();
	}

	for (auto &arg : args.data) {
		const auto &child_cols = StructVector::GetEntries(arg);
		for (auto &child_col : child_cols) {
			result_cols[offset++]->Reference(*child_col);
		}
	}
	D_ASSERT(offset == result_cols.size());

	if (args.AllConstant()) {
		result.SetVectorType(VectorType::CONSTANT_VECTOR);
	}

	result.Verify(args.size());
}

static unique_ptr<FunctionData> StructConcatBind(ClientContext &context, ScalarFunction &bound_function,
                                                 vector<unique_ptr<Expression>> &arguments) {

	// collect names and deconflict, construct return type
	if (arguments.empty()) {
		throw InvalidInputException("struct_concat: At least one argument is required");
	}

	child_list_t<LogicalType> combined_children;
	case_insensitive_set_t name_set;

	bool has_unnamed = false;

	for (idx_t arg_idx = 0; arg_idx < arguments.size(); arg_idx++) {
		const auto &arg = arguments[arg_idx];

		if (arg->return_type.id() == LogicalTypeId::UNKNOWN) {
			throw ParameterNotResolvedException();
		}

		if (arg->return_type.id() != LogicalTypeId::STRUCT) {
			throw InvalidInputException("struct_concat: Argument at position \"%d\" is not a STRUCT", arg_idx + 1);
		}

		const auto &child_types = StructType::GetChildTypes(arg->return_type);
		for (const auto &child : child_types) {
			if (!child.first.empty()) {
				auto it = name_set.find(child.first);
				if (it != name_set.end()) {
					if (*it == child.first) {
						throw InvalidInputException("struct_concat: Arguments contain duplicate STRUCT entry \"%s\"",
						                            child.first);
					}
					throw InvalidInputException(
					    "struct_concat: Arguments contain case-insensitive duplicate STRUCT entry \"%s\" and \"%s\"",
					    child.first, *it);
				}
				name_set.insert(child.first);
			} else {
				has_unnamed = true;
			}
			combined_children.push_back(child);
		}
	}

	if (has_unnamed && !name_set.empty()) {
		throw InvalidInputException("struct_concat: Cannot mix named and unnamed STRUCTs");
	}

	bound_function.return_type = LogicalType::STRUCT(combined_children);
	return nullptr;
}

static unique_ptr<BaseStatistics> StructConcatStats(ClientContext &context, FunctionStatisticsInput &input) {
	const auto &expr = input.expr;

	auto &arg_stats = input.child_stats;
	auto &arg_exprs = input.expr.children;

	auto struct_stats = StructStats::CreateUnknown(expr.return_type);
	idx_t struct_index = 0;

	for (idx_t arg_idx = 0; arg_idx < arg_exprs.size(); arg_idx++) {
		auto &arg_stat = arg_stats[arg_idx];
		auto &arg_type = arg_exprs[arg_idx]->return_type;
		for (idx_t child_idx = 0; child_idx < StructType::GetChildCount(arg_type); child_idx++) {
			auto &child_stat = StructStats::GetChildStats(arg_stat, child_idx);
			StructStats::SetChildStats(struct_stats, struct_index++, child_stat);
		}
	}
	return struct_stats.ToUnique();
}

ScalarFunction StructConcatFun::GetFunction() {
	ScalarFunction fun("struct_concat", {}, LogicalTypeId::STRUCT, StructConcatFunction, StructConcatBind, nullptr,
	                   StructConcatStats);
	fun.varargs = LogicalType::ANY;
	fun.null_handling = FunctionNullHandling::SPECIAL_HANDLING;
	return fun;
}

} // namespace duckdb
