from typing import Annotated

from haystack import component
from haystack.components.agents import Agent
from haystack.dataclasses import ChatMessage, ToolCall
from haystack.tools import tool


@tool
def get_order_status(order_id: Annotated[str, "Order ID to look up"]) -> str:
    """Return the fulfillment status for one order."""
    return f"Order {order_id} is packed and waiting for courier pickup."


@component
class DemoToolCallingChatGenerator:
    @component.output_types(replies=list[ChatMessage])
    def run(self, messages: list[ChatMessage], tools=None):
        if messages[-1].is_from("user"):
            return {
                "replies": [
                    ChatMessage.from_assistant(
                        tool_calls=[
                            ToolCall(
                                tool_name="get_order_status",
                                arguments={"order_id": "A100"},
                                id="call_1",
                            )
                        ]
                    )
                ]
            }

        tool_result = messages[-1].tool_call_result.result
        return {
            "replies": [
                ChatMessage.from_assistant(
                    f"I checked order A100: {tool_result}"
                )
            ]
        }


agent = Agent(
    chat_generator=DemoToolCallingChatGenerator(),
    tools=[get_order_status],
    system_prompt="Use tools when order status is requested.",
    max_agent_steps=4,
)

result = agent.run(
    messages=[ChatMessage.from_user("What is the status of order A100?")]
)

for message in result["messages"]:
    if message.tool_calls:
        call = message.tool_calls[0]
        print(f"tool call: {call.tool_name}({call.arguments})")
    elif message.tool_call_result:
        print(f"tool result: {message.tool_call_result.result}")

print(f"final answer: {result['last_message'].text}")