# Copyright (c) Microsoft. All rights reserved.

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

from agent_framework import (
    Agent,
    AgentResponseUpdate,
    Message,
    WorkflowEvent,
)
from agent_framework.foundry import FoundryChatClient
from agent_framework.orchestrations import MagenticBuilder, MagenticPlanReviewRequest, MagenticPlanReviewResponse
from azure.identity import AzureCliCredential
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

"""
Sample: Magentic Orchestration with Human Plan Review

This sample demonstrates how humans can review and provide feedback on plans
generated by the Magentic workflow orchestrator. When plan review is enabled,
the workflow requests human approval or revision before executing each plan.

Key concepts:
- with_plan_review(): Enables human review of generated plans
- MagenticPlanReviewRequest: The event type for plan review requests
- Human can choose to: approve the plan or provide revision feedback

Plan review options:
- approve(): Accept the proposed plan and continue execution
- revise(feedback): Provide textual feedback to modify the plan

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.
- Authentication via azure-identity. Use AzureCliCredential and run az login before executing the sample.
"""

# Keep track of the last response to format output nicely in streaming mode
last_response_id: str | None = None


async def process_event_stream(stream: AsyncIterable[WorkflowEvent]) -> dict[str, MagenticPlanReviewResponse] | None:
    """Process events from the workflow stream to capture human feedback requests."""
    global last_response_id

    requests: dict[str, MagenticPlanReviewRequest] = {}
    async for event in stream:
        if event.type == "request_info" and event.request_type is MagenticPlanReviewRequest:
            requests[event.request_id] = cast(MagenticPlanReviewRequest, event.data)

        if event.type in ("intermediate", "output"):
            data = event.data
            if isinstance(data, AgentResponseUpdate):
                rid = data.response_id
                if rid != last_response_id:
                    if last_response_id is not None:
                        print("\n")
                    print(f"{data.author_name}:", end=" ", flush=True)
                    last_response_id = rid
                print(data.text, end="", flush=True)
            else:
                # The output of the workflow comes from the orchestrator and it's a list of messages
                print("\n" + "=" * 60)
                print("DISCUSSION COMPLETE")
                print("=" * 60)
                print("Final discussion summary:")
                # To make the type checker happy, we cast event.data to the expected type
                outputs = cast(list[Message], event.data)
                for msg in outputs:
                    speaker = msg.author_name or msg.role
                    print(f"[{speaker}]: {msg.text}")

    responses: dict[str, MagenticPlanReviewResponse] = {}
    if requests:
        for request_id, request in requests.items():
            print("\n\n[Magentic Plan Review Request]")
            if request.current_progress is not None:
                print("Current Progress Ledger:")
                print(json.dumps(request.current_progress.to_dict(), indent=2))
                print()
            print(f"Proposed Plan:\n{request.plan.text}\n")
            print("Please provide your feedback (press Enter to approve):")

            reply = input("> ")  # noqa: ASYNC250
            if reply.strip() == "":
                print("Plan approved.\n")
                responses[request_id] = request.approve()
            else:
                print("Plan revised by human.\n")
                responses[request_id] = request.revise(reply)

    return responses if responses else None


async def main() -> None:
    client = FoundryChatClient(
        project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
        model=os.environ["FOUNDRY_MODEL"],
        credential=AzureCliCredential(),
    )

    researcher_agent = Agent(
        name="ResearcherAgent",
        description="Specialist in research and information gathering",
        instructions="You are a Researcher. You find information and gather facts.",
        client=client,
    )

    analyst_agent = Agent(
        name="AnalystAgent",
        description="Data analyst who processes and summarizes research findings",
        instructions="You are an Analyst. You analyze findings and create summaries.",
        client=client,
    )

    manager_agent = Agent(
        name="MagenticManager",
        description="Orchestrator that coordinates the workflow",
        instructions="You coordinate a team to complete tasks efficiently.",
        client=client,
    )

    print("\nBuilding Magentic Workflow with Human Plan Review...")

    # enable_plan_review=True: Request human input for plan review.
    # Mark participant responses as intermediate so the stream shows the
    # conversation as it unfolds while the manager's final answer remains the
    # terminal workflow output.
    workflow = MagenticBuilder(
        participants=[researcher_agent, analyst_agent],
        enable_plan_review=True,
        intermediate_output_from=[researcher_agent, analyst_agent],
        manager_agent=manager_agent,
        max_round_count=10,
        max_stall_count=1,
        max_reset_count=2,
    ).build()

    task = "Research sustainable aviation fuel technology and summarize the findings."

    print(f"\nTask: {task}")
    print("\nStarting workflow execution...")
    print("=" * 60)

    # Initiate the first run of the workflow.
    # Runs are not isolated; state is preserved across multiple calls to run.
    stream = workflow.run(task, stream=True)

    pending_responses = await process_event_stream(stream)
    while pending_responses is not None:
        # Run the workflow until there is no more human feedback to provide,
        # in which case this workflow completes.
        stream = workflow.run(stream=True, responses=pending_responses)
        pending_responses = await process_event_stream(stream)


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