# Copyright (c) Microsoft. All rights reserved.

"""Test script demonstrating middleware interception and manual function result injection.

This script demonstrates:
1. Using AzureOpenAIResponsesClient with an Agent
2. Creating a middleware that intercepts tool calls
3. Manually providing function results back to the agent by calling run again
"""

import asyncio
import os
from collections.abc import Awaitable, Callable

from agent_framework import Agent, tool
from agent_framework._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareTermination
from agent_framework.azure import AzureOpenAIResponsesClient
from azure.identity import AzureCliCredential
from dotenv import load_dotenv

load_dotenv()


class InterceptingMiddleware(FunctionMiddleware):
    """Middleware that intercepts specific tool calls and raises MiddlewareTermination.

    When a matching tool is called, this middleware short-circuits execution
    and signals that external handling is required.
    """

    def __init__(self, tool_names_to_intercept: list[str]) -> None:
        """Initialize with a list of tool names to intercept."""
        self._tool_names = set(tool_names_to_intercept)
        self.intercepted_call: FunctionInvocationContext | None = None

    async def process(
        self,
        context: FunctionInvocationContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        """Intercept matching tool calls and raise MiddlewareTermination."""
        if context.function.name not in self._tool_names:
            # Not a tool we care about, let it proceed normally
            await call_next()
            return

        print(f"[Middleware] Intercepted tool call: {context.function.name}")
        print(f"[Middleware] Arguments: {context.arguments}")

        # Store the intercepted call for external handling
        self.intercepted_call = context

        # Raise MiddlewareTermination to signal that we've handled this
        # The result will be provided externally
        raise MiddlewareTermination(result="Something were wrong")


@tool(name="get_order_status", approval_mode="never_require")
def get_order_status(order_id: str) -> str:
    """Get the status of an order by its ID.

    Args:
        order_id: The ID of the order to check.

    Returns:
        The status of the order.
    """
    # This function body won't be executed because the middleware intercepts it
    return f"Order {order_id} status: Delivered"


async def main() -> None:
    """Run the test script."""
    # Create the client
    client = AzureOpenAIResponsesClient(
        project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"],
        deployment_name=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"],
        credential=AzureCliCredential(),
    )

    # Create the intercepting middleware
    middleware = InterceptingMiddleware(tool_names_to_intercept=["get_order_status"])

    # Create an agent with the tool and middleware
    agent = Agent(
        client=client,
        name="OrderAgent",
        instructions=(
            "You are a helpful assistant that can check order statuses. "
            "When asked about an order, use the get_order_status tool."
        ),
        tools=[get_order_status],
        middleware=[middleware],
        default_options={"store": False},  # Use local history management
    )

    # Create a session to maintain conversation state
    session = agent.create_session()

    print("=" * 60)
    print("Step 1: Initial user request")
    print("=" * 60)

    # First call - the agent should invoke the tool, which gets intercepted
    user_message = "What's the status of order #12345?"
    print(f"User: {user_message}")

    response = await agent.run(user_message, session=session)

    print(f"\nAgent response (after interception): {response.text}")
    print(f"Response messages: {len(response.messages)}")

    # Check if we have an intercepted call
    if middleware.intercepted_call:
        print("\n" + "=" * 60)
        print("Step 2: Middleware intercepted the tool call")
        print("=" * 60)

        intercepted = middleware.intercepted_call
        print(f"Intercepted function: {intercepted.function.name}")
        print(f"Arguments: {intercepted.arguments}")

        # The last message from the first response
        final_response = await agent.run([response.messages[-1]], session=session, options={"tool_choice": "none"})

        print(f"Final agent response: {final_response.text}")

    else:
        print("No tool call was intercepted - the agent may have responded directly")


if __name__ == "__main__":
    asyncio.run(main())

        final_response = await agent.run([], session=session, options={"tool_choice": "none"})