/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#ifndef OMNISTREAM_OPERATOREVENTDISPATCHER_H
#define OMNISTREAM_OPERATOREVENTDISPATCHER_H

#include <unordered_map>
#include <unordered_set>
#include <memory>
#include <stdexcept>
#include <string>
#include <cassert>

// Flink includes - these would correspond to the Java imports
#include "runtime/jobgraph/tasks/TaskOperatorEventGateway.h"
#include "OperatorEvent.h"
#include "OperatorEventGateway.h"
#include "OperatorEventHandler.h"

class OperatorEventDispatcher {
public:
    /**
     * Register a listener that is notified every time an OperatorEvent is sent from the
     * OperatorCoordinator (of the operator with the given OperatorID) to this subtask.
     */
    virtual void RegisterEventHandler(const OperatorID& operatorId, std::shared_ptr<OperatorEventHandler> handler) = 0;

    /**
     * Gets the gateway through which events can be passed to the OperatorCoordinator for the
     * operator identified by the given OperatorID.
     */
    virtual std::shared_ptr<OperatorEventGateway> GetOperatorEventGateway(const OperatorID& operatorId) = 0;
};

/**
 * An implementation of the {@link OperatorEventDispatcher}.
 *
 * <p>This class is intended for single threaded use from the stream task mailbox.
 */
class OperatorEventDispatcherImpl : public OperatorEventDispatcher {
public:
    explicit OperatorEventDispatcherImpl(std::shared_ptr<TaskOperatorEventGateway> toCoordinator)
        : toCoordinator(std::move(toCoordinator)),
          handlers()
    {
    }

    void DispatchEventToHandlers(const OperatorID& operatorID, std::shared_ptr<OperatorEvent> serializedEvent);

    void RegisterEventHandler(const OperatorID& operatorId, std::shared_ptr<OperatorEventHandler> handler) override;

    std::unordered_set<OperatorID> GetRegisteredOperators();

    std::shared_ptr<OperatorEventGateway> GetOperatorEventGateway(const OperatorID& operatorId) override
    {
        return std::make_shared<OperatorEventGatewayImpl>(toCoordinator, operatorId);
    }

private:
    // ------------------------------------------------------------------------
    class OperatorEventGatewayImpl : public OperatorEventGateway {
    public:
        OperatorEventGatewayImpl(std::shared_ptr<TaskOperatorEventGateway> toCoordinator, const OperatorID& operatorId)
            : toCoordinator(toCoordinator),
              operatorId(operatorId)
        {
        }

        void SendEventToCoordinator(std::shared_ptr<OperatorEvent> event) override;

    private:
        std::shared_ptr<TaskOperatorEventGateway> toCoordinator;
        OperatorID operatorId;
    };
    std::shared_ptr<TaskOperatorEventGateway> toCoordinator;
    std::unordered_map<OperatorID, std::shared_ptr<OperatorEventHandler>> handlers;
};

#endif // OMNISTREAM_OPERATOREVENTDISPATCHER_H