#include <bit>
#include <cstdint>
#include <map>
#include <memory>
#include <utility>

#include <catch2/catch_test_macros.hpp>

#include "serif/backend/mfem/serialize/configuration.hpp"
#include "serif/backend/mfem/serialize/options.hpp"
#include "serif/backend/mfem/discretization/element_coordinates.hpp"
#include "serif/tests/solver/stellar_fixture.hpp"
#include "serif/tests/nonconforming.hpp"

namespace {
    namespace fixture = serif::tests::solver;
    namespace archive = serif::utils::serialize;
    namespace backend = serif::backend::mfem::serialize;
    using Discretization = decltype(fixture::make_discretization(std::declval<std::shared_ptr<fixture::StellarResources>>()));

    /** @brief Exercise simplex orientation independently of the moving spherical mesh generator. */
    struct FixedResources {
        ::mfem::Mesh mesh;
        ::mfem::H1_FECollection scalar_collection{fixture::Density::FunctionalSpace::order, 3};
        ::mfem::H1_FECollection displacement_collection{fixture::Displacement::FunctionalSpace::order, 3};
        ::mfem::H1_FECollection potential_collection{fixture::Potential::FunctionalSpace::order, 3};
        ::mfem::FiniteElementSpace density_space{&mesh, &scalar_collection};
        ::mfem::FiniteElementSpace displacement_space{&mesh, &displacement_collection, 3, ::mfem::Ordering::byVDIM};
        ::mfem::FiniteElementSpace potential_space{&mesh, &potential_collection};
        ::mfem::GridFunction compactification{&density_space};
        serif::discretization::backend::mfem::TermBinding<fixture::Density, fixture::Schema> density{density_space};
        serif::discretization::backend::mfem::TermBinding<fixture::Displacement, fixture::Schema> displacement{displacement_space};
        serif::discretization::backend::mfem::TermBinding<fixture::Potential, fixture::Schema> potential{potential_space};
        serif::mapping::DomainMapper mapper{{}, std::make_unique<serif::mapping::compactification::KelvinCompactification>(serif::utils::user_options::KelvinCompactificationOptions{})};

        /** @brief Finalize a supported serial mesh before constructing any term binding. */
        explicit FixedResources(::mfem::Element::Type geometry) : mesh(make_mesh(geometry)) {
            compactification = 0.0;
        }

        /** @brief Retain nontrivial support and high-order mesh coordinates in the small fixed-domain case. */
        static ::mfem::Mesh make_mesh(::mfem::Element::Type geometry) {
            auto mesh = ::mfem::Mesh::MakeCartesian3D(2, 1, 1, geometry);
            mesh.GetElement(mesh.GetNE() - 1)->SetAttribute(3);
            mesh.SetCurvature(2);
            return mesh;
        }
    };

    /** @brief Ensure restoration cannot accidentally borrow the application's original resources. */
    archive::Node saved_star(backend::SaveOptions options = {}) {
        const auto resources = fixture::make_resources(1.25, 0, 2);
        const auto discretization = fixture::make_discretization(resources, 2);
        return backend::save_discretization(discretization, options);
    }

    /** @brief Keep numerical reference choices together with the model and its native resources. */
    auto make_configuration() {
        const auto resources = fixture::make_resources(1.25, 0, 2);
        const serif::solver::normalization::NormalizationOptions normalization{
            .scheme = serif::solver::normalization::LaneEmdenApproximation{
                .n = 1.0, .mass = serif::dimensions::MassValue{1.0}, .central_density = serif::dimensions::DensityValue{1.0}}
        };
        return serif::model::make_discretization(serif::model::DiscretizationOptions{
            .model = fixture::make_model(), .finite_elements = fixture::make_discretization(resources, 2), .normalization = normalization
        });
    }

    /** @brief Compare native element coefficients after restoring their exact basis and orientation. */
    template <typename Term, typename Original, typename Restored>
    void compare_field(const Original& original, const Restored& restored) {
        const auto& before = original.template binding<Term>();
        const auto& after = restored.template binding<Term>();
        REQUIRE(before.size() == after.size());
        ::mfem::Vector coefficients(before.size());
        for (int dof = 0; dof < coefficients.Size(); ++dof) {
            coefficients[dof] = 0.125 * static_cast<double>(dof - 7);
        }
        coefficients[0] = -0.0;
        const auto before_elements = before.validated_elements();
        const auto after_elements = after.validated_elements();
        for (int element = 0; element < before.mesh().GetNE(); ++element) {
            ::mfem::Vector first, second;
            before_elements.element_values(element, coefficients, first);
            after_elements.element_values(element, coefficients, second);
            REQUIRE(first.Size() == second.Size());
            for (int local = 0; local < first.Size(); ++local) {
                CHECK(std::bit_cast<std::uint64_t>(first[local]) == std::bit_cast<std::uint64_t>(second[local]));
            }
        }
    }
}

TEST_CASE("Fixed MFEM discretizations retain simplex and tensor element orientation", "[unit][serialize][mfem_archive]") {
    namespace dependencies = serif::utils::dependencies;

    for (const ::mfem::Element::Type geometry : {::mfem::Element::HEXAHEDRON, ::mfem::Element::TETRAHEDRON}) {
        const auto resources = std::make_shared<FixedResources>(geometry);
        const auto bindings = dependencies::make_dependency_context(dependencies::borrow<fixture::Density>(resources->density),
            dependencies::borrow<fixture::Potential>(resources->potential), dependencies::borrow<fixture::Displacement>(resources->displacement));
        const auto original = serif::model::backend::mfem::make_stellar_discretization(resources, bindings, resources->compactification, resources->mapper);
        const archive::Node record = backend::save_discretization(original);
        auto restored = backend::restore_discretization<decltype(original)>(record);
        INFO((restored ? "restored" : restored.error().path + ": " + restored.error().message));
        REQUIRE(restored);
        CHECK(backend::save_discretization(*restored) == record);
        compare_field<fixture::Density>(original, *restored);
        compare_field<fixture::Potential>(original, *restored);
        compare_field<fixture::Displacement>(original, *restored);
    }
}

TEST_CASE("MFEM archives own exact meshes bases support and oriented coefficient interpretations", "[unit][serialize][mfem_archive]") {
    const auto resources = fixture::make_resources(1.25, 0, 2);
    const auto original = fixture::make_discretization(resources, 2);
    const archive::Node record = backend::save_discretization(original);
    auto restored = backend::restore_discretization<Discretization>(record);
    INFO((restored ? "restored" : restored.error().path + ": " + restored.error().message));
    REQUIRE(restored);
    CHECK(backend::save_discretization(*restored) == record);
    CHECK(&restored->template binding<fixture::Density>().mesh() != resources->mesh.mesh.get());
    CHECK(restored->deformation.reference.scalar_space.FEColl() == restored->deformation.reference.displacement_space.FEColl());
    CHECK(restored->deformation.reference.displacement_space.GetOrdering() == ::mfem::Ordering::byNODES);
    CHECK(restored->template binding<fixture::Displacement>().space().GetOrdering() == ::mfem::Ordering::byVDIM);
    CHECK(restored->deformation.model == original.deformation.model);
    compare_field<fixture::Density>(original, *restored);
    compare_field<fixture::Potential>(original, *restored);
    compare_field<fixture::Displacement>(original, *restored);
    compare_field<fixture::Constant>(original, *restored);
    compare_field<fixture::Omega>(original, *restored);

    // Nothing from the writer's resource owner survives this helper call.
    const archive::Node independent_record = saved_star();
    const auto independent = backend::restore_discretization<Discretization>(independent_record);
    REQUIRE(independent);
    CHECK(backend::save_discretization(*independent) == independent_record);
}

TEST_CASE("MFEM archives reject incompatible numbering and resource associations", "[unit][serialize][mfem_archive]") {
    archive::Node record = saved_star();
    auto& resources = record.children.at("resources");
    SECTION("Mesh contents disagree with their content identity") {
        auto& native = resources.children.at("meshes").children.at("0").children.at("native");
        std::get<std::string>(native.value).append("\n");
    }
    SECTION("A matching fingerprint does not make malformed native mesh text valid") {
        auto& mesh = resources.children.at("meshes").children.at("0");
        const std::string malformed = "This is not an MFEM mesh.\n";
        mesh.children.at("native") = archive::encode(malformed);
        mesh.children.at("identity") = archive::encode(archive::fingerprint(malformed));
        const auto restored = backend::restore_discretization<Discretization>(record);
        REQUIRE_FALSE(restored);
        CHECK(restored.error().message.starts_with("Cannot reconstruct saved MFEM resources:"));
        return;
    }
    SECTION("A collection's basis changes without changing its nominal polynomial order") {
        const auto space = archive::decode<std::uint64_t>(record.children.at("bindings").children.at("0").children.at("space"));
        REQUIRE(space);
        const auto collection_index = archive::decode<std::uint64_t>(
            resources.children.at("spaces").children.at(std::to_string(*space)).children.at("collection"));
        REQUIRE(collection_index);
        auto& collection = resources.children.at("collections").children.at(std::to_string(*collection_index));
        const auto family = archive::decode<std::string>(collection.children.at("family"));
        REQUIRE(family);
        REQUIRE(*family == "H1");
        collection.children.at("closed_basis") = archive::encode(static_cast<int>(::mfem::BasisType::ClosedUniform));
    }
    SECTION("An element's signed local numbering changes") {
        auto& numbering = resources.children.at("spaces").children.at("0").children.at("element_numbering");
        std::get<std::vector<std::int64_t>>(numbering.value)[1] = -100;
    }
    SECTION("A term is connected to a missing space") {
        record.children.at("bindings").children.at("0").children.at("space") = archive::encode(std::uint64_t{900});
    }
    SECTION("A term's reduced map is changed") {
        auto& reduction = record.children.at("bindings").children.at("0").children.at("coordinates").children.at("full_to_reduced");
        std::get<std::vector<std::int64_t>>(reduction.value)[0] = 999;
    }
    SECTION("A deformation companion is not a displacement space") {
        record.children.at("reference").children.at("displacement_space") = record.children.at("reference").children.at("scalar_space");
    }
    SECTION("Native MFEM version differs") {
        resources.children.at("mfem_version") = archive::encode(std::string{"different"});
    }
    CHECK_FALSE(backend::restore_discretization<Discretization>(record));
}

TEST_CASE("An omitted MFEM mesh must be resolved by its exact saved identity", "[unit][serialize][mfem_archive]") {
    const auto resources = fixture::make_resources(1.25, 0, 2);
    const auto original = fixture::make_discretization(resources, 2);
    const archive::Node record = backend::save_discretization(original, {.embed_mesh = false});
    CHECK_FALSE(backend::restore_discretization<Discretization>(record));
    std::map<std::string, std::string, std::less<>> meshes;
    for (const ::mfem::Mesh* mesh : {resources->mesh.mesh.get(), resources->mesh.reference_mesh.get()}) {
        const std::string text = backend::mesh_text(*mesh);
        meshes.emplace(archive::fingerprint(text), text);
    }
    auto restored = backend::restore_discretization<Discretization>(record, [&](std::string_view identity) -> archive::Result<std::string> {
        const auto found = meshes.find(identity);
        if (found == meshes.end()) {
            return std::unexpected(archive::Error{"mesh", "Unknown mesh identity."});
        }
        return found->second;
    });
    REQUIRE(restored);
    CHECK(backend::save_discretization(*restored, {.embed_mesh = false}) == record);
    CHECK_FALSE(backend::restore_discretization<Discretization>(record, [](std::string_view) -> archive::Result<std::string> { return std::string{"wrong mesh"}; }));
}

TEST_CASE("MFEM assembly options retain quadrature controls and reject unknown contraction modes", "[unit][serialize][mfem_archive]") {
    namespace quadrature = serif::discretization::quadrature;
    using Options = serif::operators::backend::mfem::VolumeOptions;

    quadrature::RuleSet rules;
    rules.set_role_control(quadrature::QuadratureRole::discretization, {.fixed_order = 9, .boost = 2});
    const Options options{.quadrature_policy = quadrature::Policy{rules}, .contraction_mode = serif::operators::backend::mfem::ContractionMode::dense};
    archive::Node record = archive::encode(options);
    auto restored = archive::decode<Options>(record);
    REQUIRE(restored);
    CHECK(archive::encode(*restored) == record);
    record.children.at("contraction") = archive::encode(500);
    CHECK_FALSE(archive::decode<Options>(record));
}

TEST_CASE("Configured MFEM models restore their full numerical prescription with independent ownership", "[unit][serialize][mfem_archive]") {
    using Configuration = decltype(make_configuration());

    archive::Node record;
    std::weak_ptr<const fixture::StellarResources> original_owner;
    {
        const auto original = make_configuration();
        original_owner = original.finite_elements.owner;
        record = archive::encode(original);
    }
    REQUIRE(original_owner.expired());
    REQUIRE(record.type == "model.configured_discretization");
    const auto restored = backend::restore_configuration<Configuration>(record);
    INFO((restored ? "restored" : restored.error().path + ": " + restored.error().message));
    REQUIRE(restored);
    const archive::Node saved_again = archive::encode(*restored);
    CHECK(saved_again.children.at("finite_elements") == record.children.at("finite_elements"));
    CHECK(saved_again.children.at("normalization") == record.children.at("normalization"));
    CHECK(saved_again.children.at("preconditioner") == record.children.at("preconditioner"));
    CHECK(saved_again.children.at("model").children.at("resolved_equations") == record.children.at("model").children.at("resolved_equations"));

    SECTION("Unknown configuration versions cannot be restored") {
        record.version = 2;
    }
    SECTION("Missing numerical choices cannot silently become defaults") {
        record.children.erase("normalization");
    }
    SECTION("Unsupported preconditioners cannot silently become no preconditioning") {
        record.children.at("preconditioner") = archive::Node{.type = "unsupported", .value = std::string{"Custom preconditioner has no serializer."}};
    }
    CHECK_FALSE(backend::restore_configuration<Configuration>(record));
}

TEST_CASE("MFEM archives refuse native parsing when the application selects abort on error", "[unit][serialize][mfem_archive]") {
    const archive::Node record = saved_star();
    struct RestoreErrorAction {
        ::mfem::ErrorAction previous{::mfem::get_error_action()};
        /** @brief Restore the application's prior native error policy even if an assertion stops the test. */
        ~RestoreErrorAction() { ::mfem::set_error_action(previous); }
    } restore_error_action;
    ::mfem::set_error_action(::mfem::MFEM_ERROR_ABORT);
    const auto restored = backend::restore_discretization<Discretization>(record);
    REQUIRE_FALSE(restored);
    CHECK(restored.error().message.find("requires MFEM_ERROR_THROW") != std::string::npos);
    CHECK(::mfem::get_error_action() == ::mfem::MFEM_ERROR_ABORT);
}

TEST_CASE("Nonconforming STROID archives retain independent coordinates after original owners expire", "[unit][serialize][mfem_archive][nonconforming]") {
    archive::Node record;
    std::weak_ptr<const void> original_owner;
    const auto restored = [&] {
        const auto resources = std::make_shared<fixture::StellarResources>(serif::tests::nonconforming::mesh(2, 0, 2));
        original_owner = resources;
        const auto original = fixture::make_discretization(resources, 2);
        REQUIRE(resources->potential_space.GetTrueVSize() < resources->potential_space.GetVSize());
        record = backend::save_discretization(original);
        auto result = backend::restore_discretization<Discretization>(record);
        if (!result) {
            INFO(result.error().path << ": " << result.error().message);
        }
        REQUIRE(result);
        compare_field<fixture::Density>(original, *result);
        compare_field<fixture::Potential>(original, *result);
        compare_field<fixture::Displacement>(original, *result);
        CHECK(&result->template binding<fixture::Potential>().mesh() != resources->mesh.mesh.get());
        return std::move(*result);
    }();
    CHECK(original_owner.expired());
    const auto& potential = restored.template binding<fixture::Potential>();
    REQUIRE(potential.space().GetTrueVSize() < potential.space().GetVSize());
    CHECK(potential.space().GetConformingProlongation() != nullptr);
    CHECK(restored.deformation.reference.logical_mesh.Nonconforming());
    CHECK(backend::save_discretization(restored) == record);
    mfem::Vector coefficients(potential.size());
    for (int index = 0; index < coefficients.Size(); ++index) {
        coefficients[index] = 0.125 * (index - 7);
    }
    const mfem::Vector local = potential.scatter(coefficients);
    const mfem::Vector independent = potential.gather(local);
    REQUIRE(independent.Size() == coefficients.Size());
    for (int index = 0; index < coefficients.Size(); ++index) {
        CHECK(independent[index] == coefficients[index]);
    }
}