# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from core.architecture.handlers.actorcommand.actor_handler_base import (
    ActorCommandHandler,
)
from core.models.internal_queue import DCMElement, InternalQueue
from core.util import set_one_of
from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity
from proto.org.apache.texera.amber.engine.architecture.rpc import (
    ControlInvocation,
    ControlRequest,
    EmptyRequest,
    AsyncRpcContext,
)
from proto.org.apache.texera.amber.engine.common import (
    Backpressure,
    DirectControlMessagePayloadV2,
)


class BackpressureHandler(ActorCommandHandler):
    cmd = Backpressure

    def __call__(
        self, command: Backpressure, input_queue: InternalQueue, *args, **kwargs
    ):
        if command.enable_backpressure:
            input_queue.disable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE)
        else:
            input_queue.enable_data(InternalQueue.DisableType.DISABLE_BY_BACKPRESSURE)
            input_queue.put(
                DCMElement(
                    tag=ChannelIdentity(
                        ActorVirtualIdentity("self"), ActorVirtualIdentity("self"), True
                    ),
                    payload=set_one_of(
                        DirectControlMessagePayloadV2,
                        ControlInvocation(
                            "NoOperation",
                            set_one_of(ControlRequest, EmptyRequest()),
                            AsyncRpcContext(),
                            -1,
                        ),
                    ),
                )
            )

        return None
