#include "gridfire/policy/stellar_policy.h"
#include "gridfire/policy/policy_abstract.h"
#include "gridfire/exceptions/error_policy.h"

#include "gridfire/engine/engine_abstract.h"
#include "gridfire/engine/engine_graph.h"
#include "gridfire/engine/views/engine_views.h"
#include "gridfire/utils/logging.h"

#include "gridfire/engine/scratchpads/blob.h"
#include "gridfire/engine/scratchpads/utils.h"
#include "gridfire/engine/scratchpads/engine_graph_scratchpad.h"
#include "gridfire/engine/scratchpads/engine_adaptive_scratchpad.h"
#include "gridfire/engine/scratchpads/engine_multiscale_scratchpad.h"

#include "fourdst/atomic/species.h"
#include "fourdst/composition/utils.h"

namespace {
    std::set<fourdst::atomic::Species> initialize_seed_species() {
        return {
            fourdst::atomic::H_1,
            fourdst::atomic::He_3,
            fourdst::atomic::He_4,
            fourdst::atomic::C_12,
            fourdst::atomic::N_14,
            fourdst::atomic::O_16,
            fourdst::atomic::Ne_20,
            fourdst::atomic::Mg_24
        };
    }
}

namespace gridfire::policy {
    MainSequencePolicy::MainSequencePolicy(const fourdst::composition::Composition& composition) : m_seed_species(initialize_seed_species()) {
        for (const auto& species : m_seed_species) {
            if (!composition.contains(species)) {
                throw exceptions::MissingSeedSpeciesError("Cannot initialize MainSequencePolicy: Required Seed species " + std::string(species.name()) + " is missing from the provided composition.");
            }
        }
        m_initializing_composition = composition;
        m_partition_function = build_partition_function();
    }

    MainSequencePolicy::MainSequencePolicy(std::vector<fourdst::atomic::Species> seed_species, const std::vector<double> &mass_fractions) {
        for (const auto& species : m_seed_species) {
            if (std::ranges::find(seed_species, species) == seed_species.end()) {
                throw exceptions::MissingSeedSpeciesError("Cannot initialize MainSequencePolicy: Required Seed species " + std::string(species.name()) + " is missing from the provided composition.");
            }
        }

        m_initializing_composition = fourdst::composition::buildCompositionFromMassFractions(seed_species, mass_fractions);

        m_partition_function = build_partition_function();
    }

     ConstructionResults MainSequencePolicy::construct() {
        m_network_stack.clear();

        m_network_stack.emplace_back(
            std::make_unique<engine::GraphEngine>(m_initializing_composition, *m_partition_function, engine::NetworkBuildDepth::ThirdOrder, engine::NetworkConstructionFlags::DEFAULT)
        );

        m_network_stack.emplace_back(
            std::make_unique<engine::MultiscalePartitioningEngineView>(*m_network_stack.back().get())
        );
        m_network_stack.emplace_back(
            std::make_unique<engine::AdaptiveEngineView>(*m_network_stack.back().get())
        );

        std::unique_ptr<engine::scratch::StateBlob> scratch_blob = get_stack_scratch_blob();
        m_status = NetworkPolicyStatus::INITIALIZED_UNVERIFIED;
        m_status = check_status(*scratch_blob);

        switch (m_status) {
            case NetworkPolicyStatus::MISSING_KEY_REACTION:
                throw exceptions::MissingKeyReactionError("MainSequencePolicy construction failed: The constructed network is missing key reactions required by the policy.");
            case NetworkPolicyStatus::MISSING_KEY_SPECIES:
                throw exceptions::MissingSeedSpeciesError("MainSequencePolicy construction failed: The constructed network is missing key seed species required by the policy.");
            case NetworkPolicyStatus::UNINITIALIZED:
                throw exceptions::PolicyError("MainSequencePolicy construction failed: The network policy is uninitialized.");
            case NetworkPolicyStatus::INITIALIZED_UNVERIFIED:
                throw exceptions::PolicyError("MainSequencePolicy construction failed: The network policy status could not be verified.");
            case NetworkPolicyStatus::INITIALIZED_VERIFIED:
                break;
        }
        return {.engine = *m_network_stack.back(), .scratch_blob = std::move(scratch_blob)};
    }

    inline std::unique_ptr<partition::PartitionFunction> MainSequencePolicy::build_partition_function() {
        using partition::BasePartitionType;
        const auto partitionFunction = partition::CompositePartitionFunction({
            BasePartitionType::RauscherThielemann,
            BasePartitionType::GroundState
        });
        return std::make_unique<partition::CompositePartitionFunction>(partitionFunction);
    }

    inline NetworkPolicyStatus MainSequencePolicy::get_status() const {
        return m_status;
    }

    const std::vector<std::unique_ptr<engine::DynamicEngine>> &MainSequencePolicy::get_engine_stack() const {
        if (m_status != NetworkPolicyStatus::INITIALIZED_VERIFIED) {
            throw exceptions::PolicyError("Cannot get engine stack from MainSequencePolicy: Policy is not initialized and verified. Call construct() first.");
        }
        return m_network_stack;
    }

    std::vector<engine::EngineTypes> MainSequencePolicy::get_engine_types_stack() const {
        return {
            engine::EngineTypes::GRAPH_ENGINE,
            engine::EngineTypes::MULTISCALE_PARTITIONING_ENGINE_VIEW,
            engine::EngineTypes::ADAPTIVE_ENGINE_VIEW
        };
    }

    const std::unique_ptr<partition::PartitionFunction>& MainSequencePolicy::get_partition_function() const {
        return m_partition_function;
    }

    std::unique_ptr<engine::scratch::StateBlob> MainSequencePolicy::get_stack_scratch_blob() const {
        if (m_network_stack.empty()) {
            throw exceptions::PolicyError("Cannot get stack scratch blob from MainSequencePolicy: Engine stack is empty. Call construct() first.");
        }
        auto blob = std::make_unique<engine::scratch::StateBlob>();
        blob->enroll<engine::scratch::GraphEngineScratchPad>();
        blob->enroll<engine::scratch::AdaptiveEngineViewScratchPad>();
        blob->enroll<engine::scratch::MultiscalePartitioningEngineViewScratchPad>();


        const engine::GraphEngine* graph_engine = dynamic_cast<engine::GraphEngine*>(m_network_stack.front().get());
        if (!graph_engine) {
            throw exceptions::PolicyError("Cannot get stack scratch blob from MainSequencePolicy: The base engine is not a GraphEngine. This indicates a serious internal inconsistency and should be reported to the GridFire developers, thank you.");
        }

        const engine::MultiscalePartitioningEngineView* multiscale_engine = dynamic_cast<engine::MultiscalePartitioningEngineView*>(m_network_stack[1].get());
        if (!multiscale_engine) {
            throw exceptions::PolicyError("Cannot get stack scratch blob from MainSequencePolicy: The middle engine is not a MultiscalePartitioningEngineView. This indicates a serious internal inconsistency and should be reported to the GridFire developers, thank you.");
        }

        const engine::AdaptiveEngineView* adaptive_engine = dynamic_cast<engine::AdaptiveEngineView*>(m_network_stack.back().get());
        if (!adaptive_engine) {
            throw exceptions::PolicyError("Cannot get stack scratch blob from MainSequencePolicy: The top engine is not an AdaptiveEngineView. This indicates a serious internal inconsistency and should be reported to the GridFire developers, thank you.");
        }


        auto* graph_engine_state = engine::scratch::get_state<engine::scratch::GraphEngineScratchPad, false>(*blob);
        graph_engine_state->initialize(*graph_engine);
        auto* multiscale_engine_state = engine::scratch::get_state<engine::scratch::MultiscalePartitioningEngineViewScratchPad, false>(*blob);
        multiscale_engine_state->initialize();
        auto* adaptive_engine_state = engine::scratch::get_state<engine::scratch::AdaptiveEngineViewScratchPad, false>(*blob);
        adaptive_engine_state->initialize(*adaptive_engine);

        return blob;
    }

    inline NetworkPolicyStatus MainSequencePolicy::check_status(engine::scratch::StateBlob& ctx) const {
        for (const auto& species : m_seed_species) {
            if (!m_initializing_composition.contains(species)) {
                return NetworkPolicyStatus::MISSING_KEY_SPECIES;
            }
        }
        const reaction::ReactionSet& baseReactions = m_network_stack.front()->getNetworkReactions(ctx);
        for (const auto& reaction : m_reaction_policy->get_reactions()) {
            const bool result = baseReactions.contains(*reaction);
            if (!result) {
                return NetworkPolicyStatus::MISSING_KEY_REACTION;
            }
        }
        return NetworkPolicyStatus::INITIALIZED_VERIFIED;
    }
}
