# Copyright (c) Microsoft. All rights reserved.

import asyncio
import datetime

from agent_framework import (
    Agent,
    agent_middleware,
    function_middleware,
    tool,
)
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

"""
Decorator MiddlewareTypes Example

This sample demonstrates how to use @agent_middleware and @function_middleware decorators
to explicitly mark middleware functions without requiring type annotations.

The framework supports the following middleware detection scenarios:

1. Both decorator and parameter type specified:
   - Validates that they match (e.g., @agent_middleware with AgentContext)
   - Throws exception if they don't match for safety

2. Only decorator specified:
   - Relies on decorator to determine middleware type
   - No type annotations needed - framework handles context types automatically

3. Only parameter type specified:
   - Uses type annotations (AgentContext, FunctionInvocationContext) for detection

4. Neither decorator nor parameter type specified:
   - Throws exception requiring either decorator or type annotation
   - Prevents ambiguous middleware that can't be properly classified

Key benefits of decorator approach:
- No type annotations needed (simpler syntax)
- Explicit middleware type declaration
- Clear intent in code
- Prevents type mismatches
"""


# 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_current_time() -> str:
    """Get the current time."""
    return f"Current time is {datetime.datetime.now().strftime('%H:%M:%S')}"


@agent_middleware  # Decorator marks this as agent middleware - no type annotations needed
async def simple_agent_middleware(context, call_next):  # type: ignore - parameters intentionally untyped to demonstrate decorator functionality
    """Agent middleware that runs before and after agent execution."""
    print("[Agent MiddlewareTypes] Before agent execution")
    await call_next()
    print("[Agent MiddlewareTypes] After agent execution")


@function_middleware  # Decorator marks this as function middleware - no type annotations needed
async def simple_function_middleware(context, call_next):  # type: ignore - parameters intentionally untyped to demonstrate decorator functionality
    """Function middleware that runs before and after function calls."""
    print(f"[Function MiddlewareTypes] Before calling: {context.function.name}")  # type: ignore
    await call_next()
    print(f"[Function MiddlewareTypes] After calling: {context.function.name}")  # type: ignore


async def main() -> None:
    """Example demonstrating decorator-based middleware."""
    print("=== Decorator MiddlewareTypes Example ===")

    # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
    # authentication option.
    async with (
        AzureCliCredential() as credential,
        Agent(
            client=FoundryChatClient(credential=credential),
            name="TimeAgent",
            instructions="You are a helpful time assistant. Call get_current_time when asked about time.",
            tools=get_current_time,
            middleware=[simple_agent_middleware, simple_function_middleware],
        ) as agent,
    ):
        query = "What time is it?"
        print(f"User: {query}")
        result = await agent.run(query)
        print(f"Agent: {result.text if result.text else 'No response'}")


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