#include <memory>
#include <tuple>
#include <type_traits>

#include <catch2/catch_test_macros.hpp>

#include "serif/discretization/blocks/fields.hpp"
#include "serif/equations/concepts.hpp"
#include "serif/operators/contributions.hpp"

namespace {
    /// Record exactly when an equation's owned numerical coefficient is destroyed.
    struct TrackedCoefficient {
        double value;
        std::shared_ptr<int> destructions;

        /// Construct the tracked value without an intermediate temporary.
        TrackedCoefficient(double coefficient, std::shared_ptr<int> count) : value(coefficient), destructions(std::move(count)) {}
        /// Release the observable coefficient lifetime.
        ~TrackedCoefficient() { ++*destructions; }
    };

    /// A move-only leaf prevents traversal from concealing lifetime mistakes through equation copies.
    struct TrackedEquation {
        using Test = serif::discretization::blocks::density::mass;
        using Density = serif::discretization::forms::Operand<Test>;
        using Operands = serif::utils::types::TypeList<Density>;
        static constexpr std::string_view name = "tracked equation";

        std::unique_ptr<TrackedCoefficient> slope;

        /// Own a numerical coefficient whose destruction remains observable outside this equation.
        TrackedEquation(double coefficient, const std::shared_ptr<int>& count) : slope(std::make_unique<TrackedCoefficient>(coefficient, count)) {}

        /// Evaluate a linear residual using the retained coefficient.
        [[nodiscard]] serif::equations::PointLinearization<Operands> evaluate(const serif::equations::EquationPoint<Operands>& point) const {
            serif::equations::PointLinearization<Operands> result;
            result.residual = slope->value * point.fields.get<Density>();
            result.field_derivatives.get<Density>() = slope->value;
            return result;
        }
    };

    /// Generate a nonempty move-only contribution tuple by value, as mixed gravity generates its leaves.
    struct GeneratedEquation {
        using Test = TrackedEquation::Test;
        static constexpr std::string_view name = "generated contribution";
        double slope;
        std::shared_ptr<int> destructions;

        /// The caller must retain this tuple if it wants to retain references to its equation.
        [[nodiscard]] auto contributions() const { return std::tuple{TrackedEquation{slope, destructions}}; }
    };

    static_assert(serif::equations::Equation<GeneratedEquation>);
    static_assert(!std::is_copy_constructible_v<TrackedEquation>);
}

TEST_CASE("Contribution bindings retain generated tuples and borrowed move-only leaves", "[equations][contributions][unit]") {
    using namespace serif::equations;

    const auto destructions = std::make_shared<int>(0);
    auto owner = std::make_shared<const GeneratedEquation>(3.0, destructions);
    {
        const auto leaves = serif::operators::bind_contributions(owner);
        STATIC_CHECK(std::tuple_size_v<decltype(leaves)> == 1);
        owner.reset();

        // The generated tuple remains alive after its parent equation is released.
        REQUIRE(*destructions == 0);
        CHECK(std::get<0>(leaves).coefficient == 1.0);
        EquationPoint<TrackedEquation::Operands> point;
        point.fields.get<TrackedEquation::Density>() = 4.0;
        CHECK(std::get<0>(leaves).equation->evaluate(point).residual == 12.0);
    }
    CHECK(*destructions == 1);
}