#pragma once

#include <grove/IIO.h>
#include <grove/JsonDataNode.h>
#include <nlohmann/json.hpp>

#include <string>
#include <vector>
#include <queue>
#include <map>
#include <algorithm>

namespace celuna::tests {

using json = nlohmann::json;

/**
 * @brief Mock implementation of grove::IIO for testing
 *
 * Captures published messages and allows injecting incoming messages.
 */
class MockIO : public grove::IIO {
public:
    // ========================================================================
    // IIO Interface Implementation
    // ========================================================================

    void publish(const std::string& topic, std::unique_ptr<grove::IDataNode> data) override;

    void subscribe(const std::string& topicPattern, const grove::SubscriptionConfig& config = {}) override {
        // Mock: just record subscription
        m_subscriptions.push_back(topicPattern);
    }

    void subscribeLowFreq(const std::string& topicPattern, const grove::SubscriptionConfig& config = {}) override {
        // Mock: same as subscribe
        m_subscriptions.push_back(topicPattern);
    }

    int hasMessages() const override {
        return static_cast<int>(m_incomingMessages.size());
    }

    grove::Message pullMessage() override;

    grove::IOHealth getHealth() const override {
        return grove::IOHealth{
            .queueSize = static_cast<int>(m_incomingMessages.size()),
            .maxQueueSize = 1000,
            .dropping = false,
            .averageProcessingRate = 100.0f,
            .droppedMessageCount = 0
        };
    }

    grove::IOType getType() const override {
        return grove::IOType::INTRA;
    }

    // ========================================================================
    // Test Helpers - Message Injection
    // ========================================================================

    /**
     * @brief Inject a message to be received by the module under test
     */
    void injectMessage(const std::string& topic, const json& data);

    /**
     * @brief Inject multiple messages at once
     */
    void injectMessages(const std::vector<std::pair<std::string, json>>& messages);

    // ========================================================================
    // Test Helpers - Verification
    // ========================================================================

    /**
     * @brief Check if a message was published to a specific topic
     */
    bool wasPublished(const std::string& topic) const;

    /**
     * @brief Get the last message published to a topic
     */
    json getLastPublished(const std::string& topic) const;

    /**
     * @brief Get all messages published to a topic
     */
    std::vector<json> getAllPublished(const std::string& topic) const;

    /**
     * @brief Count messages published to a topic
     */
    size_t countPublished(const std::string& topic) const;

    /**
     * @brief Get all published messages (topic -> data pairs)
     */
    const std::vector<std::pair<std::string, json>>& getPublishedMessages() const {
        return m_publishedMessages;
    }

    /**
     * @brief Clear all captured and pending messages
     */
    void clear();

    /**
     * @brief Clear only published messages (keep incoming queue)
     */
    void clearPublished();

    // ========================================================================
    // Test State
    // ========================================================================

    /// All messages published by the module under test
    std::vector<std::pair<std::string, json>> m_publishedMessages;

    /// Messages waiting to be received by the module
    std::queue<grove::Message> m_incomingMessages;

    /// Subscribed topic patterns (for verification)
    std::vector<std::string> m_subscriptions;
};

} // namespace celuna::tests
