import json

from strands import Agent
from strands.experimental.hooks.multiagent import BeforeNodeCallEvent
from strands.hooks import HookProvider
from strands.interrupt import Interrupt
from strands.multiagent import Swarm
from strands.multiagent.base import Status


class ApprovalHook(HookProvider):
    def register_hooks(self, registry):
        registry.add_callback(BeforeNodeCallEvent, self.approve)

    def approve(self, event):
        if event.node_id == "info":
            return

        response = event.interrupt("my_interrupt", reason=f"{event.node_id} needs approval")
        if response != "APPROVE":
            event.cancel_node = "node rejected"


info_agent = Agent(name="info")
weather_agent = Agent(name="weather")
swarm = Swarm([info_agent, weather_agent], hooks=[ApprovalHook()])

multiagent_result = swarm("What is the weather?")
while multiagent_result.status == Status.INTERRUPTED:
    response = []
    for interrupt in multiagent_result.interrupts:
        if interrupt.name == "my_interrupt":
            response = input(f"{interrupt.reason} (y/N): ")
            responses.append(
                {
                    "interruptResponse": {
                         "interruptId": interrupt.id,
                         "response": response,
                    },
                },
             )

    multiagent_result = swarm(responses)


weather_result = multiagent_result.results["weather"]
weather_message = json.dumps(weather_result.result.message).lower()
print(weather_message)

import json

from strands import Agent, tool
from strands.interrupt import Interrupt
from strands.multiagent import Swarm
from strands.multiagent.base import Status


@tool(context=True)
def user_name_tool(tool_context: ToolContext) -> str:
    response = tool_context.interrupt("my_interrupt", reason="need user name")
    return response


user_agent = Agent(name="user", tools=[user_name_tool])
swarm = Swarm([user_agent])

multiagent_result = swarm("Who is the current user?")
while multiagent_result.status == Status.INTERRUPTED:
    response = []
    for interrupt in multiagent_result.interrupts:
        if interrupt.name == "my_interrupt":
            response = input(f"{interrupt.reason}: ")
            responses.append(
                {
                    "interruptResponse": {
                         "interruptId": interrupt.id,
                         "response": response,
                    },
                },
             )

    multiagent_result = swarm(responses)


user_result = multiagent_result.results["user"]
user_message = json.dumps(user_result.result.message).lower()
print(user_message)