#include "RootNode.hpp"

#include "MayaFlux/Nodes/Generators/Generator.hpp"

namespace MayaFlux::Nodes {

RootNode::RootNode(ProcessingToken token, uint32_t channel)
    : m_is_processing(false)
    , m_pending_count(0)
    , m_channel(channel)
    , m_skip_state_management(false)
    , m_token(token)
{
}

void RootNode::register_node(const std::shared_ptr<Node>& node)
{
    if (!node)
        return;

    for (auto& pending_op : m_pending_ops) {
        bool expected = false;
        if (pending_op.active.compare_exchange_strong(
                expected, true,
                std::memory_order_acquire,
                std::memory_order_relaxed)) {
            pending_op.node = node;
            pending_op.is_addition = true;
            atomic_remove_flag(node->m_state, NodeState::ACTIVE);
            atomic_add_flag(node->m_state, NodeState::INACTIVE);
            m_pending_count.fetch_add(1, std::memory_order_relaxed);
            return;
        }
    }

    while (m_is_processing.load(std::memory_order_acquire))
        m_is_processing.wait(true, std::memory_order_acquire);

    if (m_Nodes.end() == std::ranges::find(m_Nodes, node)) {
        m_Nodes.push_back(node);
        atomic_remove_flag(node->m_state, NodeState::INACTIVE);
        atomic_add_flag(node->m_state, NodeState::ACTIVE);
    }
}

void RootNode::unregister_node(const std::shared_ptr<Node>& node)
{
    if (!node)
        return;

    atomic_add_flag(node->m_state, NodeState::PENDING_REMOVAL);

    for (auto& pending_op : m_pending_ops) {
        bool expected = false;
        if (pending_op.active.compare_exchange_strong(
                expected, true,
                std::memory_order_acquire,
                std::memory_order_relaxed)) {
            pending_op.node = node;
            pending_op.is_addition = false;
            m_pending_count.fetch_add(1, std::memory_order_relaxed);
            return;
        }
    }

    while (m_is_processing.load(std::memory_order_acquire))
        m_is_processing.wait(true, std::memory_order_acquire);

    auto it = m_Nodes.begin();
    while (it != m_Nodes.end()) {
        if ((*it).get() == node.get()) {
            it = m_Nodes.erase(it);
            break;
        }
        ++it;
    }

    node->reset_processed_state();

    uint32_t flag = node->m_state.load();
    flag &= ~static_cast<uint32_t>(NodeState::PENDING_REMOVAL);
    flag &= ~static_cast<uint32_t>(NodeState::ACTIVE);
    flag |= static_cast<uint32_t>(NodeState::INACTIVE);
    atomic_set_flag_strong(node->m_state, static_cast<NodeState>(flag));
}

bool RootNode::preprocess()
{
    if (m_skip_state_management)
        return true;

    bool expected = false;
    if (m_request_terminate.load(std::memory_order_acquire)) {
        return false;
    }

    if (!m_is_processing.compare_exchange_strong(expected, true,
            std::memory_order_acquire, std::memory_order_relaxed)) {
        return false;
    }

    if (m_pending_count.load(std::memory_order_relaxed) > 0) {
        process_pending_operations();
    }

    return true;
}

double RootNode::process_sample()
{
    if (!preprocess())
        return 0.;

    auto sample = 0.;

    for (auto& node : m_Nodes) {
        if (!node)
            continue;

        uint32_t state = node->m_state.load();
        double node_output = 0.0;

        if (!(state & NodeState::PROCESSED)) {
            auto generator = std::dynamic_pointer_cast<Nodes::Generator::Generator>(node);
            if (generator && generator->should_mock_process()) {
                generator->process_sample();
            } else {
                node_output = node->process_sample();
            }
            atomic_add_flag(node->m_state, NodeState::PROCESSED);
        } else {
            node_output = node->get_last_output();
        }

        if (node->needs_channel_routing()) {
            node_output *= node->get_routing_state().amount[m_channel];
        }

        sample += node_output;
    }

    postprocess();

    return sample;
}

void RootNode::process_frame()
{
    if (!preprocess())
        return;

    for (auto& node : m_Nodes) {
        uint32_t state = node->m_state.load();
        if (!(state & NodeState::PROCESSED)) {
            node->process_sample();
            atomic_add_flag(node->m_state, NodeState::PROCESSED);
        }
    }

    postprocess();
}

void RootNode::postprocess()
{
    if (m_skip_state_management)
        return;

    for (auto& node : m_Nodes) {
        node->request_reset_from_channel(m_channel);
    }

    if (m_pending_count.load(std::memory_order_relaxed) > 0) {
        process_pending_operations();
    }

    m_is_processing.store(false, std::memory_order_release);
    m_is_processing.notify_all();

    if (m_request_terminate.load(std::memory_order_acquire)) {
        process_pending_operations();
    }
}

std::vector<double> RootNode::process_batch(uint32_t num_samples)
{
    std::vector<double> output(num_samples);

    for (unsigned int i = 0; i < num_samples; i++) {
        output[i] = process_sample();
    }
    return output;
}

void RootNode::process_batch_frame(uint32_t num_frames)
{
    for (uint32_t i = 0; i < num_frames; i++) {
        process_frame();
    }
}

void RootNode::process_pending_operations()
{
    for (auto& pending_op : m_pending_ops) {
        if (!pending_op.active.load(std::memory_order_acquire))
            continue;

        auto& op = pending_op;

        if (op.is_addition) {
            if (m_Nodes.end() == std::ranges::find(m_Nodes, op.node)) {
                m_Nodes.push_back(op.node);
                uint32_t state = op.node->m_state.load();
                state &= ~static_cast<uint32_t>(NodeState::INACTIVE);
                state |= static_cast<uint32_t>(NodeState::ACTIVE);
                atomic_set_flag_strong(op.node->m_state, static_cast<NodeState>(state));
            }
        } else {
            auto it = m_Nodes.begin();
            while (it != m_Nodes.end()) {
                if ((*it).get() == op.node.get()) {
                    it = m_Nodes.erase(it);
                    break;
                }
                ++it;
            }
            op.node->reset_processed_state();
            uint32_t state = op.node->m_state.load();
            state &= ~static_cast<uint32_t>(NodeState::PENDING_REMOVAL);
            state &= ~static_cast<uint32_t>(NodeState::ACTIVE);
            state |= static_cast<uint32_t>(NodeState::INACTIVE);
            atomic_set_flag_strong(op.node->m_state, static_cast<NodeState>(state));
        }

        op.node.reset();
        op.active.store(false, std::memory_order_release);
        m_pending_count.fetch_sub(1, std::memory_order_relaxed);
    }
}

void RootNode::terminate_all_nodes()
{
    m_request_terminate.store(true, std::memory_order_release);

    m_is_processing.store(false, std::memory_order_release);

    for (auto& node : m_Nodes) {
        unregister_node(node);
    }

    process_pending_operations();
}

}
