# Copyright (c) Microsoft. All rights reserved.

import asyncio
import json
import os
from typing import Annotated, Any, cast

from agent_framework import Message, tool
from agent_framework.azure import AzureOpenAIResponsesClient
from agent_framework.orchestrations import SequentialBuilder
from azure.identity import AzureCliCredential
from dotenv import load_dotenv
from pydantic import Field

# Load environment variables from .env file
load_dotenv()

"""
Sample: Workflow kwargs Flow to @tool Tools

This sample demonstrates how to flow custom context (skill data, user tokens, etc.)
through any workflow pattern to @tool functions using the **kwargs pattern.

Key Concepts:
- Pass custom context as kwargs when invoking workflow.run()
- kwargs are stored in State and passed to all agent invocations
- @tool functions receive kwargs via **kwargs parameter
- Works with Sequential, Concurrent, GroupChat, Handoff, and Magentic patterns

Prerequisites:
- AZURE_AI_PROJECT_ENDPOINT must be your Azure AI Foundry Agent Service (V2) project endpoint.
- Environment variables configured
"""


# Define tools that accept custom context via **kwargs
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_user_data(
    query: Annotated[str, Field(description="What user data to retrieve")],
    **kwargs: Any,
) -> str:
    """Retrieve user-specific data based on the authenticated context."""
    user_token = kwargs.get("user_token", {})
    user_name = user_token.get("user_name", "anonymous")
    access_level = user_token.get("access_level", "none")

    print(f"\n[get_user_data] Received kwargs keys: {list(kwargs.keys())}")
    print(f"[get_user_data] User: {user_name}")
    print(f"[get_user_data] Access level: {access_level}")

    return f"Retrieved data for user {user_name} with {access_level} access: {query}"


@tool(approval_mode="never_require")
def call_api(
    endpoint_name: Annotated[str, Field(description="Name of the API endpoint to call")],
    **kwargs: Any,
) -> str:
    """Call an API using the configured endpoints from custom_data."""
    custom_data = kwargs.get("custom_data", {})
    api_config = custom_data.get("api_config", {})

    base_url = api_config.get("base_url", "unknown")
    endpoints = api_config.get("endpoints", {})

    print(f"\n[call_api] Received kwargs keys: {list(kwargs.keys())}")
    print(f"[call_api] Base URL: {base_url}")
    print(f"[call_api] Available endpoints: {list(endpoints.keys())}")

    if endpoint_name in endpoints:
        return f"Called {base_url}{endpoints[endpoint_name]} successfully"
    return f"Endpoint '{endpoint_name}' not found in configuration"


async def main() -> None:
    print("=" * 70)
    print("Workflow kwargs Flow Demo (SequentialBuilder)")
    print("=" * 70)

    # Create chat client
    client = AzureOpenAIResponsesClient(
        project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"],
        deployment_name=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"],
        credential=AzureCliCredential(),
    )

    # Create agent with tools that use kwargs
    agent = client.as_agent(
        name="assistant",
        instructions=(
            "You are a helpful assistant. Use the available tools to help users. "
            "When asked about user data, use get_user_data. "
            "When asked to call an API, use call_api."
        ),
        tools=[get_user_data, call_api],
    )

    # Build a simple sequential workflow
    workflow = SequentialBuilder(participants=[agent]).build()

    # Define custom context that will flow to tools via kwargs
    custom_data = {
        "api_config": {
            "base_url": "https://api.example.com",
            "endpoints": {
                "users": "/v1/users",
                "orders": "/v1/orders",
                "products": "/v1/products",
            },
        },
    }

    user_token = {
        "user_name": "bob@contoso.com",
        "access_level": "admin",
    }

    print("\nCustom Data being passed:")
    print(json.dumps(custom_data, indent=2))
    print(f"\nUser: {user_token['user_name']}")
    print("\n" + "-" * 70)
    print("Workflow Execution (watch for [tool_name] logs showing kwargs received):")
    print("-" * 70)

    # Run workflow with kwargs - these will flow through to tools
    async for event in workflow.run(
        "Please get my user data and then call the users API endpoint.",
        additional_function_arguments={"custom_data": custom_data, "user_token": user_token},
        stream=True,
    ):
        if event.type == "output":
            output_data = cast(list[Message], event.data)
            if isinstance(output_data, list):
                for item in output_data:
                    if isinstance(item, Message) and item.text:
                        print(f"\n[Final Answer]: {item.text}")

    print("\n" + "=" * 70)
    print("Sample Complete")
    print("=" * 70)


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