# Copyright (c) Microsoft. All rights reserved.

"""
Handoff Workflow with Code Interpreter File Generation Sample

This sample demonstrates retrieving file IDs from code interpreter output
in a handoff workflow context. A triage agent routes to a code specialist
that generates a text file, and we verify the file_id is captured correctly
from the streaming workflow events.

Verifies GitHub issue #2718: files generated by code interpreter in
HandoffBuilder workflows can be properly retrieved.

Prerequisites:
    - FOUNDRY_PROJECT_ENDPOINT must be your Azure AI Foundry Agent Service (V2) project endpoint.
    - FOUNDRY_MODEL must be set to your Azure OpenAI model deployment name.
    - `az login` (Azure CLI authentication)
"""

import asyncio
import os
from collections.abc import AsyncIterable
from typing import cast

from agent_framework import (
    Agent,
    AgentResponseUpdate,
    Message,
    WorkflowEvent,
    WorkflowRunState,
)
from agent_framework.foundry import FoundryChatClient
from agent_framework.orchestrations import HandoffAgentUserRequest, HandoffBuilder
from azure.identity import AzureCliCredential
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()


async def _drain(stream: AsyncIterable[WorkflowEvent]) -> list[WorkflowEvent]:
    """Collect all events from an async stream."""
    return [event async for event in stream]


def _handle_events(events: list[WorkflowEvent]) -> tuple[list[WorkflowEvent[HandoffAgentUserRequest]], list[str]]:
    """Process workflow events and extract file IDs and pending requests.

    Returns:
        Tuple of (pending_requests, file_ids_found)
    """

    requests: list[WorkflowEvent[HandoffAgentUserRequest]] = []
    file_ids: list[str] = []

    for event in events:
        if event.type == "handoff_sent":
            print(f"\n[Handoff from {event.data.source} to {event.data.target} initiated.]")
        elif event.type == "status" and event.state in {
            WorkflowRunState.IDLE,
            WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
        }:
            print(f"[status] {event.state}")
        elif event.type == "request_info" and isinstance(event.data, HandoffAgentUserRequest):
            requests.append(cast(WorkflowEvent[HandoffAgentUserRequest], event))
        elif event.type == "output":
            data = event.data
            if isinstance(data, AgentResponseUpdate):
                for content in data.contents:
                    if content.type == "hosted_file":
                        file_ids.append(content.file_id)  # type: ignore
                        print(f"[Found HostedFileContent: file_id={content.file_id}]")
                    elif content.type == "text" and content.annotations:
                        for annotation in content.annotations:
                            file_id = annotation["file_id"]  # type: ignore
                            file_ids.append(file_id)
                            print(f"[Found file annotation: file_id={file_id}]")
            elif isinstance(data, list):
                conversation = cast(list[Message], data)
                if isinstance(conversation, list):
                    print("\n=== Final Conversation Snapshot ===")
                    for message in conversation:
                        speaker = message.author_name or message.role
                        print(f"- {speaker}: {message.text or [content.type for content in message.contents]}")
                    print("===================================")

    return requests, file_ids


async def main() -> None:
    """Run a simple handoff workflow with code interpreter file generation."""
    print("=== Handoff Workflow with Code Interpreter File Generation ===\n")

    client = FoundryChatClient(
        project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
        model=os.environ["FOUNDRY_MODEL"],
        credential=AzureCliCredential(),
    )

    triage = Agent(
        client=client,
        name="triage_agent",
        instructions=(
            "You are a triage agent. Route code-related requests to the code_specialist. "
            "When the user asks to create or generate files, hand off to code_specialist "
            "by calling handoff_to_code_specialist."
        ),
        require_per_service_call_history_persistence=True,
    )

    code_interpreter_tool = client.get_code_interpreter_tool()

    code_specialist = Agent(
        client=client,
        name="code_specialist",
        instructions=(
            "You are a Python code specialist. Use the code interpreter to execute Python code "
            "and create files when requested. Always save files to /mnt/data/ directory."
        ),
        tools=[code_interpreter_tool],
        require_per_service_call_history_persistence=True,
    )

    workflow = (
        HandoffBuilder(
            termination_condition=lambda conv: sum(1 for msg in conv if msg.role == "user") >= 2,
        )
        .participants([triage, code_specialist])
        .with_start_agent(triage)
        .build()
    )

    user_inputs = [
        "Please create a text file called hello.txt with 'Hello from handoff workflow!' inside it.",
        "exit",
    ]
    input_index = 0
    all_file_ids: list[str] = []

    print(f"User: {user_inputs[0]}")
    events = await _drain(workflow.run(user_inputs[0], stream=True))
    requests, file_ids = _handle_events(events)
    all_file_ids.extend(file_ids)
    input_index += 1

    while requests:
        request = requests[0]
        if input_index >= len(user_inputs):
            break
        user_input = user_inputs[input_index]
        print(f"\nUser: {user_input}")

        responses = {request.request_id: HandoffAgentUserRequest.create_response(user_input)}
        events = await _drain(workflow.run(stream=True, responses=responses))
        requests, file_ids = _handle_events(events)
        all_file_ids.extend(file_ids)
        input_index += 1

    print("\n" + "=" * 50)
    if all_file_ids:
        print(f"SUCCESS: Found {len(all_file_ids)} file ID(s) in handoff workflow:")
        for fid in all_file_ids:
            print(f"  - {fid}")
    else:
        print("WARNING: No file IDs captured from the handoff workflow.")
    print("=" * 50)

    """
    Sample Output:

    User: Please create a text file called hello.txt with 'Hello from handoff workflow!' inside it.
    [Found HostedFileContent: file_id=assistant-JT1sA...]

    === Conversation So Far ===
    - user: Please create a text file called hello.txt with 'Hello from handoff workflow!' inside it.
    - triage_agent: I am handing off your request to create the text file "hello.txt" with the specified content to the code specialist. They will assist you shortly.
    - code_specialist: The file "hello.txt" has been created with the content "Hello from handoff workflow!". You can download it using the link below:

    [hello.txt](sandbox:/mnt/data/hello.txt)
    ===========================

    [status] IDLE_WITH_PENDING_REQUESTS

    User: exit
    [status] IDLE

    ==================================================
    SUCCESS: Found 1 file ID(s) in handoff workflow:
    - assistant-JT1sA...
    ==================================================
    """  # noqa: E501


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