module;

#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <memory>
#include <optional>
#include <span>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>

#include <mfem.hpp>
#include <mpi.h>

export module mean_field:solver.stellar_structure;

export import :operators.stellar_equilibrium_problem;
export import :solver.stellar_equilibrium_types;

export namespace mean_field::solver {
    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem> class StellarEquilibriumEvaluationReport;
}

namespace mean_field::solver::detail {
    enum class StellarViewCertification : std::uint8_t { unavailable, checkpoint, structure };

    template <typename Problem> struct StellarStructureStorage final {
        using ProblemType = std::remove_cvref_t<Problem>;

        StellarStructureStorage(
            std::unique_ptr<ProblemType> ownedProblem,
            std::unique_ptr<mfem::Vector> acceptedPhysicalState,
            std::unique_ptr<physics::RigidRotation> prescribedRotation
        )
            : problem(std::move(ownedProblem)),
              physicalState(std::move(acceptedPhysicalState)),
              rotation(std::move(prescribedRotation)) {
        }

        std::unique_ptr<ProblemType> problem;
        std::unique_ptr<mfem::Vector> physicalState;
        std::unique_ptr<physics::RigidRotation> rotation;
        std::uint64_t viewGeneration{0};
        StellarViewCertification certification{StellarViewCertification::unavailable};
    };

    template <typename Problem>
    [[nodiscard]] bool IsCurrentView(
        const std::weak_ptr<const StellarStructureStorage<Problem>> &candidate,
        const std::uint64_t generation,
        const bool requireConverged
    ) noexcept {
        const auto storage = candidate.lock();
        if (storage == nullptr || storage->problem == nullptr || storage->physicalState == nullptr ||
            storage->viewGeneration != generation) {
            return false;
        }
        if (requireConverged) {
            return storage->certification == StellarViewCertification::structure;
        }
        return storage->certification == StellarViewCertification::checkpoint ||
               storage->certification == StellarViewCertification::structure;
    }

    template <typename Problem>
    [[nodiscard]] std::shared_ptr<const StellarStructureStorage<Problem>> RequireCurrentView(
        const std::weak_ptr<const StellarStructureStorage<Problem>> &candidate,
        const std::uint64_t generation,
        const bool requireConverged
    ) {
        auto storage = candidate.lock();
        if (storage == nullptr || storage->problem == nullptr || storage->physicalState == nullptr ||
            storage->viewGeneration != generation ||
            (requireConverged && storage->certification != StellarViewCertification::structure) ||
            (!requireConverged && storage->certification != StellarViewCertification::checkpoint &&
             storage->certification != StellarViewCertification::structure)) {
            throw std::logic_error("The stellar structure view is stale or is not certified for this result.");
        }
        return storage;
    }

    template <typename Vector> [[nodiscard]] std::span<const mfem::real_t> ReadOnlySpan(const Vector &values) noexcept {
        return {values.GetData(), static_cast<std::size_t>(values.Size())};
    }

    template <typename Problem> struct StellarEvaluationReportAccess;
} // namespace mean_field::solver::detail

export namespace mean_field::equilibrium {
    /*
     * These are the future owning, self-contained values. There is no public
     * construction path until deep capture and its MPI-independent storage
     * schema are implemented.
     */
    template <DiscretizedStellarEquilibriumProblem Problem> class StellarStructure final {
    public:
        using ProblemType                                     = std::remove_cvref_t<Problem>;

        StellarStructure(const StellarStructure &)            = delete;
        StellarStructure &operator=(const StellarStructure &) = delete;
        StellarStructure(StellarStructure &&) noexcept        = default;
        StellarStructure &operator=(StellarStructure &&)      = delete;
        ~StellarStructure()                                   = default;

    private:
        StellarStructure() = default;
    };

    template <DiscretizedStellarEquilibriumProblem Problem> class StellarCheckpoint final {
    public:
        using ProblemType                                       = std::remove_cvref_t<Problem>;

        StellarCheckpoint(const StellarCheckpoint &)            = delete;
        StellarCheckpoint &operator=(const StellarCheckpoint &) = delete;
        StellarCheckpoint(StellarCheckpoint &&) noexcept        = default;
        StellarCheckpoint &operator=(StellarCheckpoint &&)      = delete;
        ~StellarCheckpoint()                                    = default;

    private:
        StellarCheckpoint() = default;
    };

    /*
     * Result views weakly observe context-owned storage. valid() remains safe
     * after that context is destroyed. References and spans extracted from a
     * valid view remain borrowed: the context must outlive their use, and the
     * next evaluate() call invalidates them along with their originating view.
     */
    template <DiscretizedStellarEquilibriumProblem Problem> class StellarStructureView final {
    public:
        using ProblemType = std::remove_cvref_t<Problem>;
        using ModelType   = typename ProblemType::ModelType;

        [[nodiscard]] bool valid() const noexcept {
            return solver::detail::IsCurrentView(m_storage, m_generation, true);
        }

        [[nodiscard]] const ModelType &model() const & {
            const auto storage = RequireStorage();
            return storage->problem->GetStellarModel();
        }

        [[nodiscard]] const ModelType &model() const && = delete;

        [[nodiscard]] MPI_Comm communicator() const & {
            const auto storage = RequireStorage();
            return storage->problem->GetCommunicator();
        }

        [[nodiscard]] MPI_Comm communicator() const && = delete;

        [[nodiscard]] std::span<const mfem::real_t> state() const & {
            const auto storage = RequireStorage();
            return solver::detail::ReadOnlySpan(*storage->physicalState);
        }

        [[nodiscard]] std::span<const mfem::real_t> state() const && = delete;

        [[nodiscard]] std::span<const operators::RootBlockDescriptor> stateDescriptors() const & {
            const auto storage = RequireStorage();
            return storage->problem->GetManifest().valueBlocks();
        }

        [[nodiscard]] std::span<const operators::RootBlockDescriptor> stateDescriptors() const && = delete;

        template <typename Term>
            requires requires(
                const typename ProblemType::ManifestType &manifest,
                const mfem::Vector &physicalState,
                const Term &term
            ) { manifest.stateView(physicalState).block(term); }
        [[nodiscard]] std::span<const mfem::real_t> stateBlock(const Term &term) const & {
            const auto storage = RequireStorage();
            const auto block   = storage->problem->GetManifest().stateView(*storage->physicalState).block(term);
            return solver::detail::ReadOnlySpan(block);
        }

        template <typename Term> [[nodiscard]] std::span<const mfem::real_t> stateBlock(const Term &) const && = delete;

        [[nodiscard]] std::optional<physics::RigidRotation> prescribedRotation() const & {
            const auto storage = RequireStorage();
            if (storage->rotation == nullptr) {
                return std::nullopt;
            }
            return *storage->rotation;
        }

        [[nodiscard]] std::optional<physics::RigidRotation> prescribedRotation() const && = delete;

        [[nodiscard]] physics::RigidRotation rotation() const & {
            const auto storage = RequireStorage();
            return storage->problem->GetPreparedOperator().GetRotation();
        }

        [[nodiscard]] physics::RigidRotation rotation() const && = delete;

        [[nodiscard]] StellarStructure<ProblemType> capture() const {
            (void)RequireStorage();
            throw std::logic_error("Capturing a self-contained StellarStructure is not implemented.");
        }

    private:
        template <DiscretizedStellarEquilibriumProblem> friend class solver::StellarEquilibriumEvaluationReport;
        using Storage = solver::detail::StellarStructureStorage<ProblemType>;

        StellarStructureView(
            std::weak_ptr<const Storage> storage,
            const std::uint64_t generation
        ) noexcept
            : m_storage(std::move(storage)),
              m_generation(generation) {
        }

        [[nodiscard]] std::shared_ptr<const Storage> RequireStorage() const {
            return solver::detail::RequireCurrentView(m_storage, m_generation, true);
        }

        std::weak_ptr<const Storage> m_storage;
        std::uint64_t m_generation;
    };

    template <DiscretizedStellarEquilibriumProblem Problem> class StellarCheckpointView final {
    public:
        using ProblemType = std::remove_cvref_t<Problem>;
        using ModelType   = typename ProblemType::ModelType;

        [[nodiscard]] bool valid() const noexcept {
            return solver::detail::IsCurrentView(m_storage, m_generation, false);
        }

        [[nodiscard]] const ModelType &model() const & {
            const auto storage = RequireStorage();
            return storage->problem->GetStellarModel();
        }

        [[nodiscard]] const ModelType &model() const && = delete;

        [[nodiscard]] MPI_Comm communicator() const & {
            const auto storage = RequireStorage();
            return storage->problem->GetCommunicator();
        }

        [[nodiscard]] MPI_Comm communicator() const && = delete;

        [[nodiscard]] std::span<const mfem::real_t> state() const & {
            const auto storage = RequireStorage();
            return solver::detail::ReadOnlySpan(*storage->physicalState);
        }

        [[nodiscard]] std::span<const mfem::real_t> state() const && = delete;

        [[nodiscard]] std::span<const operators::RootBlockDescriptor> stateDescriptors() const & {
            const auto storage = RequireStorage();
            return storage->problem->GetManifest().valueBlocks();
        }

        [[nodiscard]] std::span<const operators::RootBlockDescriptor> stateDescriptors() const && = delete;

        template <typename Term>
            requires requires(
                const typename ProblemType::ManifestType &manifest,
                const mfem::Vector &physicalState,
                const Term &term
            ) { manifest.stateView(physicalState).block(term); }
        [[nodiscard]] std::span<const mfem::real_t> stateBlock(const Term &term) const & {
            const auto storage = RequireStorage();
            const auto block   = storage->problem->GetManifest().stateView(*storage->physicalState).block(term);
            return solver::detail::ReadOnlySpan(block);
        }

        template <typename Term> [[nodiscard]] std::span<const mfem::real_t> stateBlock(const Term &) const && = delete;

        [[nodiscard]] std::optional<physics::RigidRotation> prescribedRotation() const & {
            const auto storage = RequireStorage();
            if (storage->rotation == nullptr) {
                return std::nullopt;
            }
            return *storage->rotation;
        }

        [[nodiscard]] std::optional<physics::RigidRotation> prescribedRotation() const && = delete;

        [[nodiscard]] physics::RigidRotation rotation() const & {
            const auto storage = RequireStorage();
            return storage->problem->GetPreparedOperator().GetRotation();
        }

        [[nodiscard]] physics::RigidRotation rotation() const && = delete;

        [[nodiscard]] StellarCheckpoint<ProblemType> capture() const {
            (void)RequireStorage();
            throw std::logic_error("Capturing a self-contained StellarCheckpoint is not implemented.");
        }

    private:
        template <DiscretizedStellarEquilibriumProblem> friend class solver::StellarEquilibriumEvaluationReport;
        using Storage = solver::detail::StellarStructureStorage<ProblemType>;

        StellarCheckpointView(
            std::weak_ptr<const Storage> storage,
            const std::uint64_t generation
        ) noexcept
            : m_storage(std::move(storage)),
              m_generation(generation) {
        }

        [[nodiscard]] std::shared_ptr<const Storage> RequireStorage() const {
            return solver::detail::RequireCurrentView(m_storage, m_generation, false);
        }

        std::weak_ptr<const Storage> m_storage;
        std::uint64_t m_generation;
    };

    template <DiscretizedStellarEquilibriumProblem Problem>
    [[noreturn]] void serialize(
        const StellarStructure<Problem> &,
        const std::filesystem::path &
    ) {
        throw std::logic_error("Serializing a StellarStructure is not implemented.");
    }

    template <DiscretizedStellarEquilibriumProblem Problem>
    [[noreturn]] void serialize(
        const StellarStructureView<Problem> &view,
        const std::filesystem::path &
    ) {
        (void)view.state();
        throw std::logic_error("Serializing a StellarStructureView is not implemented.");
    }

    template <DiscretizedStellarEquilibriumProblem Problem>
    [[noreturn]] void serialize(
        const StellarCheckpoint<Problem> &,
        const std::filesystem::path &
    ) {
        throw std::logic_error("Serializing a StellarCheckpoint is not implemented.");
    }

    template <DiscretizedStellarEquilibriumProblem Problem>
    [[noreturn]] void serialize(
        const StellarCheckpointView<Problem> &view,
        const std::filesystem::path &
    ) {
        (void)view.state();
        throw std::logic_error("Serializing a StellarCheckpointView is not implemented.");
    }
} // namespace mean_field::equilibrium

export namespace mean_field::solver {
    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
    class StellarEquilibriumEvaluationReport final {
    public:
        using ProblemType    = std::remove_cvref_t<Problem>;
        using StructureView  = equilibrium::StellarStructureView<ProblemType>;
        using CheckpointView = equilibrium::StellarCheckpointView<ProblemType>;

        StellarEquilibriumEvaluationReport(const StellarEquilibriumEvaluationReport &)                = default;
        StellarEquilibriumEvaluationReport &operator=(const StellarEquilibriumEvaluationReport &)     = default;
        StellarEquilibriumEvaluationReport(StellarEquilibriumEvaluationReport &&) noexcept            = default;
        StellarEquilibriumEvaluationReport &operator=(StellarEquilibriumEvaluationReport &&) noexcept = default;
        ~StellarEquilibriumEvaluationReport()                                                         = default;

        [[nodiscard]] bool converged() const noexcept {
            return m_converged;
        }

        [[nodiscard]] const StellarEquilibriumEvaluationDiagnostics &diagnostics() const & noexcept {
            return m_diagnostics;
        }

        [[nodiscard]] const StellarEquilibriumEvaluationDiagnostics &diagnostics() const && = delete;

        [[nodiscard]] int completedNonlinearIterations() const noexcept {
            return m_diagnostics.acceptedNonlinearIterations;
        }

        [[nodiscard]] double initialResidualNorm() const noexcept {
            return m_diagnostics.initialResidualNorm;
        }

        [[nodiscard]] double finalResidualNorm() const noexcept {
            return m_diagnostics.finalResidualNorm;
        }

        [[nodiscard]] const StellarEquilibriumFailureReport &failure() const & {
            if (!m_failure.has_value()) {
                throw std::logic_error("A converged stellar-equilibrium report has no failure record.");
            }
            return *m_failure;
        }

        [[nodiscard]] const StellarEquilibriumFailureReport &failure() const && = delete;

        [[nodiscard]] StructureView structureView() const {
            if (!m_converged) {
                throw std::logic_error("A failed stellar-equilibrium report cannot certify a structure view.");
            }
            StructureView view{m_storage, m_generation};
            if (!view.valid()) {
                throw std::logic_error("The stellar-equilibrium structure view has been invalidated.");
            }
            return view;
        }

        [[nodiscard]] CheckpointView checkpointView() const {
            CheckpointView view{m_storage, m_generation};
            if (!view.valid()) {
                throw std::logic_error("The stellar-equilibrium checkpoint view has been invalidated.");
            }
            return view;
        }

        [[nodiscard]] CheckpointView lastAcceptedCheckpointView() const {
            return checkpointView();
        }

    private:
        friend struct detail::StellarEvaluationReportAccess<ProblemType>;
        using Storage = detail::StellarStructureStorage<ProblemType>;

        StellarEquilibriumEvaluationReport(
            const bool converged,
            StellarEquilibriumEvaluationDiagnostics diagnostics,
            std::optional<StellarEquilibriumFailureReport> failure,
            std::weak_ptr<const Storage> storage,
            const std::uint64_t generation
        )
            : m_converged(converged),
              m_diagnostics(std::move(diagnostics)),
              m_failure(std::move(failure)),
              m_storage(std::move(storage)),
              m_generation(generation) {
        }

        bool m_converged;
        StellarEquilibriumEvaluationDiagnostics m_diagnostics;
        std::optional<StellarEquilibriumFailureReport> m_failure;
        std::weak_ptr<const Storage> m_storage;
        std::uint64_t m_generation;
    };
} // namespace mean_field::solver

namespace mean_field::solver::detail {
    template <typename Problem> struct StellarEvaluationReportAccess final {
        using ProblemType = std::remove_cvref_t<Problem>;
        using Report      = StellarEquilibriumEvaluationReport<ProblemType>;
        using Storage     = StellarStructureStorage<ProblemType>;

        [[nodiscard]] static Report Success(
            const std::shared_ptr<Storage> &storage,
            StellarEquilibriumEvaluationDiagnostics diagnostics
        ) {
            if (storage == nullptr) {
                throw std::invalid_argument("A stellar-equilibrium report requires owned result storage.");
            }
            storage->certification = StellarViewCertification::structure;
            return Report{true, std::move(diagnostics), std::nullopt, storage, storage->viewGeneration};
        }

        [[nodiscard]] static Report Failure(
            const std::shared_ptr<Storage> &storage,
            StellarEquilibriumEvaluationDiagnostics diagnostics,
            const StellarEquilibriumFailureReason reason,
            std::string message
        ) {
            if (storage == nullptr) {
                throw std::invalid_argument("A stellar-equilibrium report requires owned result storage.");
            }
            storage->certification = StellarViewCertification::checkpoint;
            StellarEquilibriumFailureReport failure{
                .reason                       = reason,
                .message                      = std::move(message),
                .completedNonlinearIterations = diagnostics.acceptedNonlinearIterations,
                .initialResidualNorm          = diagnostics.initialResidualNorm,
                .finalResidualNorm            = diagnostics.finalResidualNorm
            };
            return Report{
                false, std::move(diagnostics), std::optional<StellarEquilibriumFailureReport>{std::move(failure)},
                storage, storage->viewGeneration
            };
        }
    };
} // namespace mean_field::solver::detail