#include "fesa/assembly/sparse_assembler.h"

#include <cstddef>
#include <functional>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "fesa/analysis/analysis_model.h"
#include "fesa/assembly/parallel_for.h"
#include "fesa/elements/element_factory.h"
#include "fesa/fem/dof_manager.h"
#include "fesa/model/domain.h"

#include "assembly/ordered_contribution_internal.h"

namespace fesa {
namespace {

Status AssemblyFailure(const std::string& code, const SourceLocation& location,
                       const std::string& identity,
                       const std::string& message) {
  return Status::Failure(
      FailureCategory::kModel,
      {{Severity::kError, code, location, "*ELEMENT", identity, message}});
}

bool SameLayout(const ElementDofLayout& left, const ElementDofLayout& right) {
  return SourceIdentityEquals(left.source_id, right.source_id) &&
         left.node_indices == right.node_indices &&
         left.components_per_node == right.components_per_node;
}

}  // namespace

Result<SparseMatrix> SparseAssembler::Assemble(
    const ElementView& elements, const DofManager& dofs,
    const ParallelFor& parallel_for) {
  std::vector<std::vector<CooContribution>> local_buffers(elements.size());
  std::vector<std::optional<Status>> local_failures(elements.size());
  parallel_for.Execute(elements.size(), [&](const std::size_t element_order) {
    const Element& element = elements[element_order].get();
    auto stiffness = element.ComputeStiffness();
    if (!stiffness.HasValue()) {
      local_failures[element_order] = stiffness.GetStatus();
      return;
    }

    const auto& contribution = stiffness.Value();
    const auto& layout = element.DofLayout();
    if (!SameLayout(contribution.layout, layout)) {
      local_failures[element_order] = AssemblyFailure(
          "invalid-assembly-layout", {}, layout.source_id.source_label_text,
          "Element stiffness layout must match its declared runtime layout.");
      return;
    }
    auto scatter = dofs.ElementScatter(layout);
    if (!scatter.HasValue()) {
      local_failures[element_order] = scatter.GetStatus();
      return;
    }
    const std::size_t local_dof_count = scatter.Value().size();
    if (contribution.values.Rows() != local_dof_count ||
        contribution.values.Columns() != local_dof_count ||
        local_dof_count >
            (std::numeric_limits<std::size_t>::max)() / local_dof_count) {
      local_failures[element_order] = AssemblyFailure(
          "invalid-assembly-dimensions", {}, layout.source_id.source_label_text,
          "Element stiffness dimensions must match its declared DOF layout.");
      return;
    }

    auto& buffer = local_buffers[element_order];
    buffer.reserve(local_dof_count * local_dof_count);
    assembly_internal::AppendLocalMatrix(
        contribution.values, scatter.Value(), element_order, buffer);
  });

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

  std::size_t contribution_count = 0U;
  for (const auto& buffer : local_buffers) {
    if (buffer.size() >
        (std::numeric_limits<std::size_t>::max)() - contribution_count) {
      return Result<SparseMatrix>::Failure(AssemblyFailure(
          "invalid-assembly-dimensions", {}, std::to_string(elements.size()),
          "Element contribution storage exceeds the addressable range."));
    }
    contribution_count += buffer.size();
  }
  std::vector<CooContribution> contributions;
  contributions.reserve(contribution_count);
  // Flatten only after all workers finish. This fixed element order remains
  // independent of serial, reverse, or TBB task completion order.
  for (const auto& buffer : local_buffers) {
    contributions.insert(contributions.end(), buffer.begin(), buffer.end());
  }
  return SparseMatrix::FromCoo(dofs.FullDofCount(), dofs.FullDofCount(),
                               std::move(contributions),
                               dofs.GetSparsePattern());
}

Result<SparseMatrix> SparseAssembler::AssembleStiffness(
    const AnalysisModel& model, const DofManager& dofs,
    const ParallelFor& parallel_for) {
  const Domain& domain = model.GetDomain();
  std::vector<std::unique_ptr<Element>> owned_elements;
  ElementView elements;
  owned_elements.reserve(model.ActiveElements().size());
  elements.reserve(model.ActiveElements().size());

  const ElementFactory factory;
  for (const EntityIndex element_index : model.ActiveElements()) {
    if (element_index >= domain.Elements().Size()) {
      return Result<SparseMatrix>::Failure(
          AssemblyFailure("invalid-assembly-element", {domain.SourcePath(), 0U},
                          std::to_string(element_index),
                          "Active element index is outside the Domain."));
    }
    auto candidate = factory.Create(domain.Elements()[element_index], domain);
    if (!candidate.HasValue()) {
      return Result<SparseMatrix>::Failure(candidate.GetStatus());
    }
    owned_elements.push_back(std::move(candidate.Value()));
    elements.push_back(std::cref(*owned_elements.back()));
  }
  return Assemble(elements, dofs, parallel_for);
}

}  // namespace fesa