#include "fesa/assembly/sparse_assembler.hpp"

#include "fesa/analysis/analysis_model.hpp"
#include "fesa/assembly/parallel_for.hpp"
#include "fesa/elements/euler_beam_3d.hpp"
#include "fesa/fem/dof_manager.hpp"

#include <array>
#include <limits>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace fesa {
namespace {

constexpr std::size_t kDofsPerNode = 6U;
constexpr std::size_t kElementDofCount = 12U;
constexpr std::size_t kContributionCount =
    kElementDofCount * kElementDofCount;

using ElementBuffer = std::array<CooContribution, kContributionCount>;

Result<SparseMatrix> assemblyFailure(
    const std::string& code,
    const SourceLocation& location,
    const std::string& identity,
    const std::string& message) {
    return Result<SparseMatrix>::failure(Status::failure(
        FailureCategory::model,
        {{Severity::error,
          code,
          location,
          "*ELEMENT",
          identity,
          message}}));
}

} // namespace

Result<SparseMatrix> SparseAssembler::assembleStiffness(
    const AnalysisModel& model,
    const DofManager& dofs,
    const ParallelFor& parallelFor) {
    const Domain& domain = model.domain();
    if (domain.nodes().size() >
        (std::numeric_limits<std::size_t>::max)() / kDofsPerNode ||
        dofs.fullDofCount() != domain.nodes().size() * kDofsPerNode) {
        return assemblyFailure(
            "invalid-assembly-dimensions",
            {domain.sourcePath(), 0U},
            std::to_string(dofs.fullDofCount()),
            "DofManager dimensions do not match the active model nodes.");
    }
    if (model.activeElements().size() >
        (std::numeric_limits<std::size_t>::max)() / kContributionCount) {
        return assemblyFailure(
            "invalid-assembly-dimensions",
            {domain.sourcePath(), 0U},
            std::to_string(model.activeElements().size()),
            "Element contribution storage exceeds the addressable range.");
    }

    std::vector<std::array<std::size_t, kElementDofCount>> scatters;
    scatters.reserve(model.activeElements().size());
    for (const EntityIndex elementIndex : model.activeElements()) {
        if (elementIndex >= domain.elements().size()) {
            return assemblyFailure(
                "invalid-assembly-element",
                {domain.sourcePath(), 0U},
                std::to_string(elementIndex),
                "Active element index is outside the Domain.");
        }
        const auto& element = domain.elements()[elementIndex];
        if (element.nodeIndices[0U] >= domain.nodes().size() ||
            element.nodeIndices[1U] >= domain.nodes().size() ||
            element.materialIndex >= domain.materials().size() ||
            element.sectionIndex >= domain.sections().size()) {
            return assemblyFailure(
                "invalid-assembly-element",
                element.location,
                element.sourceId.sourceLabelText,
                "Element references an entity outside the Domain.");
        }

        std::array<std::size_t, kElementDofCount> scatter{};
        try {
            scatter = dofs.elementScatter(elementIndex);
        } catch (const std::out_of_range&) {
            return assemblyFailure(
                "invalid-assembly-scatter",
                element.location,
                element.sourceId.sourceLabelText,
                "DofManager does not contain the active element scatter.");
        }
        for (std::size_t endpoint = 0U; endpoint < 2U; ++endpoint) {
            for (std::size_t component = 0U;
                 component < kDofsPerNode;
                 ++component) {
                const std::size_t local = endpoint * kDofsPerNode + component;
                const std::size_t expected =
                    static_cast<std::size_t>(element.nodeIndices[endpoint]) *
                        kDofsPerNode +
                    component;
                if (scatter[local] != expected ||
                    scatter[local] >= dofs.fullDofCount()) {
                    return assemblyFailure(
                        "invalid-assembly-scatter",
                        element.location,
                        element.sourceId.sourceLabelText,
                        "Element scatter does not match the active model topology.");
                }
            }
        }
        scatters.push_back(scatter);
    }

    std::vector<ElementBuffer> localBuffers(model.activeElements().size());
    std::vector<std::optional<Status>> localFailures(
        model.activeElements().size());
    parallelFor.execute(
        model.activeElements().size(),
        [&](const std::size_t elementOrder) {
            const EntityIndex elementIndex = model.activeElements()[elementOrder];
            const auto& definition = domain.elements()[elementIndex];
            const auto beam = EulerBeam3D::create(
                domain.nodes()[definition.nodeIndices[0U]],
                domain.nodes()[definition.nodeIndices[1U]],
                domain.sections()[definition.sectionIndex],
                domain.materials()[definition.materialIndex]);
            if (!beam.hasValue()) {
                localFailures[elementOrder] = beam.status();
                return;
            }

            const Matrix stiffness = beam.value().globalStiffness();
            auto& buffer = localBuffers[elementOrder];
            const auto& scatter = scatters[elementOrder];
            for (std::size_t localRow = 0U;
                 localRow < kElementDofCount;
                 ++localRow) {
                for (std::size_t localColumn = 0U;
                     localColumn < kElementDofCount;
                     ++localColumn) {
                    const std::size_t localOrder =
                        localRow * kElementDofCount + localColumn;
                    buffer[localOrder] = {
                        scatter[localRow],
                        scatter[localColumn],
                        stiffness(localRow, localColumn),
                        elementOrder,
                        localOrder};
                }
            }
        });

    for (std::size_t elementOrder = 0U;
         elementOrder < localFailures.size();
         ++elementOrder) {
        if (localFailures[elementOrder]) {
            return Result<SparseMatrix>::failure(
                *localFailures[elementOrder]);
        }
    }

    std::vector<CooContribution> contributions;
    contributions.reserve(
        localBuffers.size() * kContributionCount);
    // Flatten only after all workers complete; workers never share CSR state.
    for (const auto& buffer : localBuffers) {
        contributions.insert(
            contributions.end(), buffer.begin(), buffer.end());
    }
    return SparseMatrix::fromCoo(
        dofs.fullDofCount(),
        dofs.fullDofCount(),
        std::move(contributions),
        dofs.sparsePattern());
}

} // namespace fesa
