# Copyright (c) Microsoft. All rights reserved.
# type: ignore
"""
Multi-Agent Travel Planning Workflow Evaluation with Multiple Response Tracking

This sample demonstrates a multi-agent travel planning workflow using the Azure AI Client that:
1. Processes travel queries through 7 specialized agents
2. Tracks MULTIPLE response and conversation IDs per agent for evaluation
3. Uses the new Prompt Agents API (V2)
4. Captures complete interaction sequences including multiple invocations
5. Aggregates findings through a travel planning coordinator

WORKFLOW STRUCTURE (7 agents):
- Travel Agent Executor → Hotel Search, Flight Search, Activity Search (fan-out)
- Hotel Search Executor → Booking Information Aggregation Executor
- Flight Search Executor → Booking Information Aggregation Executor
- Booking Information Aggregation Executor → Booking Confirmation Executor
- Booking Confirmation Executor → Booking Payment Executor
- Booking Information Aggregation, Booking Payment, Activity Search → Travel Planning Coordinator (ResearchLead) for final aggregation (fan-in)

Agents:
1. Travel Agent - Main coordinator (no tools to avoid thread conflicts)
2. Hotel Search - Searches hotels with tools
3. Flight Search - Searches flights with tools
4. Activity Search - Searches activities with tools
5. Booking Information Aggregation - Aggregates hotel & flight booking info
6. Booking Confirmation - Confirms bookings with tools
7. Booking Payment - Processes payments with tools
"""

import asyncio
import os
from collections import defaultdict

from _tools import (
    check_flight_availability,
    check_hotel_availability,
    confirm_booking,
    get_flight_details,
    get_hotel_details,
    process_payment,
    search_activities,
    search_flights,
    # Travel planning tools
    search_hotels,
    validate_payment_method,
)
from agent_framework import (
    Agent,
    AgentExecutorResponse,
    AgentResponseUpdate,
    Executor,
    Message,
    WorkflowBuilder,
    WorkflowContext,
    WorkflowEvent,
    executor,
    handler,
)
from agent_framework.foundry import FoundryChatClient
from azure.ai.projects.aio import AIProjectClient
from azure.identity.aio import DefaultAzureCredential
from dotenv import load_dotenv
from typing_extensions import Never

load_dotenv()


@executor(id="start_executor")
async def start_executor(input: str, ctx: WorkflowContext[list[Message]]) -> None:
    """Initiates the workflow by sending the user query to all specialized agents."""
    await ctx.send_message([Message("user", [input])])


class ResearchLead(Executor):
    """Aggregates and summarizes travel planning findings from all specialized agents."""

    def __init__(self, client: FoundryChatClient, id: str = "travel-planning-coordinator"):
        # Use default_options to persist conversation history for evaluation.
        self.agent = Agent(
            client=client,
            id="travel-planning-coordinator",
            instructions=(
                "You are the final coordinator. You will receive responses from multiple agents: "
                "booking-info-aggregation-agent (hotel/flight options), booking-payment-agent (payment confirmation), "
                "and activity-search-agent (activities). "
                "Review each agent's response, then create a comprehensive travel itinerary organized by: "
                "1. Flights 2. Hotels 3. Activities 4. Booking confirmations 5. Payment details. "
                "Clearly indicate which information came from which agent. Do not use tools."
            ),
            name="travel-planning-coordinator",
        )
        super().__init__(id=id)

    @handler
    async def fan_in_handle(self, responses: list[AgentExecutorResponse], ctx: WorkflowContext[Never, str]) -> None:
        user_query = responses[0].full_conversation[0].text

        # Extract findings from all agent responses
        agent_findings = self._extract_agent_findings(responses)
        summary_text = (
            "\n".join(agent_findings) if agent_findings else "No specific findings were provided by the agents."
        )

        # Generate comprehensive travel plan summary
        messages = [
            Message(
                role="system",
                contents=[
                    "You are a travel planning coordinator. Summarize findings from multiple specialized travel agents and provide a clear, comprehensive travel plan based on the user's query."
                ],
            ),
            Message(
                role="user",
                contents=[
                    f"Original query: {user_query}\n\nFindings from specialized travel agents:\n{summary_text}\n\nPlease provide a comprehensive travel plan based on these findings."
                ],
            ),
        ]

        try:
            final_response = await self.agent.run(messages)
            output_text = (
                final_response.messages[-1].text
                if final_response.messages and final_response.messages[-1].text
                else f"Based on the available findings, here's your travel plan for '{user_query}': {summary_text}"
            )
        except Exception:
            output_text = f"Based on the available findings, here's your travel plan for '{user_query}': {summary_text}"

        await ctx.yield_output(output_text)

    def _extract_agent_findings(self, responses: list[AgentExecutorResponse]) -> list[str]:
        """Extract findings from agent responses."""
        agent_findings = []

        for response in responses:
            findings = []
            if response.agent_response and response.agent_response.messages:
                for msg in response.agent_response.messages:
                    if msg.role == "assistant" and msg.text and msg.text.strip():
                        findings.append(msg.text.strip())

            if findings:
                combined_findings = " ".join(findings)
                agent_findings.append(f"[{response.executor_id}]: {combined_findings}")

        return agent_findings


async def run_workflow_with_response_tracking(
    query: str, client: FoundryChatClient | None = None, model: str | None = None
) -> dict:
    """Run multi-agent workflow and track conversation IDs, response IDs, and interaction sequence.

    Args:
        query: The user query to process through the multi-agent workflow
        client: Optional FoundryChatClient instance
        model: Optional model for the workflow agents

    Returns:
        Dictionary containing interaction sequence, conversation/response IDs, and conversation analysis
    """
    if client is None:
        try:
            async with DefaultAzureCredential() as credential:
                project_client = AIProjectClient(
                    endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
                    credential=credential,
                )

                async with project_client:
                    client = FoundryChatClient(project_client=project_client, model=model)
                    return await _run_workflow_with_client(query, client)
        except Exception as e:
            print(f"Error during workflow execution: {e}")
            raise
    else:
        return await _run_workflow_with_client(query, client)


async def _run_workflow_with_client(query: str, client: FoundryChatClient) -> dict:
    """Execute workflow with given client and track all interactions."""

    # Initialize tracking variables - use lists to track multiple responses per agent
    conversation_ids: dict[str, list[str]] = defaultdict(list)
    response_ids: dict[str, list[str]] = defaultdict(list)

    # Create workflow components using a single shared client
    workflow, agent_map = await _create_workflow(client)

    def track_ids(event: WorkflowEvent) -> WorkflowEvent:
        """Transform hook that tracks response/conversation IDs from AgentResponseUpdate events."""
        if event.type == "output" and isinstance(event.data, AgentResponseUpdate):
            _track_agent_ids(event, event.executor_id, response_ids, conversation_ids)
        return event

    # Process workflow events using a transform hook for ID tracking
    stream = workflow.run(query, stream=True).with_transform_hook(track_ids)
    result = await stream.get_final_response()

    workflow_output = result.get_outputs()[-1] if result.get_outputs() else None
    if workflow_output:
        print(f"\nWorkflow Output: {workflow_output}\n")

    return {
        "conversation_ids": dict(conversation_ids),
        "response_ids": dict(response_ids),
        "output": workflow_output,
        "query": query,
    }


async def _create_workflow(client: FoundryChatClient):
    """Create the multi-agent travel planning workflow with specialized agents.

    Uses a single shared FoundryChatClient for all agents.
    """

    final_coordinator = ResearchLead(client=client, id="final-coordinator")

    # Agent 1: Travel Request Handler (initial coordinator)
    travel_request_handler = Agent(
        client=client,
        id="travel-request-handler",
        instructions=(
            "You receive user travel queries and relay them to specialized agents. Extract key information: destination, dates, budget, and preferences. Pass this information forward clearly to the next agents."
        ),
        name="travel-request-handler",
    )

    # Agent 2: Hotel Search Executor
    hotel_search_agent = Agent(
        client=client,
        id="hotel-search-agent",
        instructions=(
            "You are a hotel search specialist. Your task is ONLY to search for and provide hotel information. Use search_hotels to find options, get_hotel_details for specifics, and check_availability to verify rooms. Output format: List hotel names, prices per night, total cost for the stay, locations, ratings, amenities, and addresses. IMPORTANT: Only provide hotel information without additional commentary."
        ),
        name="hotel-search-agent",
        tools=[search_hotels, get_hotel_details, check_hotel_availability],
    )

    # Agent 3: Flight Search Executor
    flight_search_agent = Agent(
        client=client,
        id="flight-search-agent",
        instructions=(
            "You are a flight search specialist. Your task is ONLY to search for and provide flight information. Use search_flights to find options, get_flight_details for specifics, and check_availability for seats. Output format: List flight numbers, airlines, departure/arrival times, prices, durations, and cabin class. IMPORTANT: Only provide flight information without additional commentary."
        ),
        name="flight-search-agent",
        tools=[search_flights, get_flight_details, check_flight_availability],
    )

    # Agent 4: Activity Search Executor
    activity_search_agent = Agent(
        client=client,
        id="activity-search-agent",
        instructions=(
            "You are an activities specialist. Your task is ONLY to search for and provide activity information. Use search_activities to find options for activities. Output format: List activity names, descriptions, prices, durations, ratings, and categories. IMPORTANT: Only provide activity information without additional commentary."
        ),
        name="activity-search-agent",
        tools=[search_activities],
    )

    # Agent 5: Booking Confirmation Executor
    booking_confirmation_agent = Agent(
        client=client,
        id="booking-confirmation-agent",
        instructions=(
            "You confirm bookings. Use check_hotel_availability and check_flight_availability to verify slots, then confirm_booking to finalize. Provide ONLY: confirmation numbers, booking references, and confirmation status."
        ),
        name="booking-confirmation-agent",
        tools=[confirm_booking, check_hotel_availability, check_flight_availability],
    )

    # Agent 6: Booking Payment Executor
    booking_payment_agent = Agent(
        client=client,
        id="booking-payment-agent",
        instructions=(
            "You process payments. Use validate_payment_method to verify payment, then process_payment to complete transactions. Provide ONLY: payment confirmation status, transaction IDs, and payment amounts."
        ),
        name="booking-payment-agent",
        tools=[process_payment, validate_payment_method],
    )

    # Agent 7: Booking Information Aggregation Executor
    booking_info_aggregation_agent = Agent(
        client=client,
        id="booking-info-aggregation-agent",
        instructions=(
            "You aggregate hotel and flight search results. Receive options from search agents and organize them. Provide: top 2-3 hotel options with prices and top 2-3 flight options with prices in a structured format."
        ),
        name="booking-info-aggregation-agent",
    )

    # Build workflow with logical booking flow:
    # 1. start_executor → travel_request_handler
    # 2. travel_request_handler → hotel_search, flight_search, activity_search (fan-out)
    # 3. hotel_search → booking_info_aggregation
    # 4. flight_search → booking_info_aggregation
    # 5. booking_info_aggregation → booking_confirmation
    # 6. booking_confirmation → booking_payment
    # 7. booking_info_aggregation, booking_payment, activity_search → final_coordinator (final aggregation, fan-in)

    workflow = (
        WorkflowBuilder(name="Travel Planning Workflow", start_executor=start_executor)
        .add_edge(start_executor, travel_request_handler)
        .add_fan_out_edges(travel_request_handler, [hotel_search_agent, flight_search_agent, activity_search_agent])
        .add_edge(hotel_search_agent, booking_info_aggregation_agent)
        .add_edge(flight_search_agent, booking_info_aggregation_agent)
        .add_edge(booking_info_aggregation_agent, booking_confirmation_agent)
        .add_edge(booking_confirmation_agent, booking_payment_agent)
        .add_fan_in_edges(
            [booking_info_aggregation_agent, booking_payment_agent, activity_search_agent], final_coordinator
        )
        .build()
    )

    # Return workflow and agent map for thread ID extraction
    agent_map = {
        "travel_request_handler": travel_request_handler,
        "hotel-search-agent": hotel_search_agent,
        "flight-search-agent": flight_search_agent,
        "activity-search-agent": activity_search_agent,
        "booking-confirmation-agent": booking_confirmation_agent,
        "booking-payment-agent": booking_payment_agent,
        "booking-info-aggregation-agent": booking_info_aggregation_agent,
        "final-coordinator": final_coordinator.agent,
    }

    return workflow, agent_map


def _track_agent_ids(event, agent, response_ids, conversation_ids):
    """Track agent response and conversation IDs - supporting multiple responses per agent."""
    update = event.data

    # response_id is directly on AgentResponseUpdate
    if update.response_id and update.response_id not in response_ids[agent]:
        response_ids[agent].append(update.response_id)

    # conversation_id is on the underlying ChatResponseUpdate (raw_representation)
    raw = update.raw_representation
    if (
        raw
        and hasattr(raw, "conversation_id")
        and raw.conversation_id
        and raw.conversation_id not in conversation_ids[agent]
    ):
        conversation_ids[agent].append(raw.conversation_id)


async def create_and_run_workflow(model: str | None = None):
    """Run the workflow evaluation and display results.

    Args:
        model: Optional model for the workflow agents

    Returns:
        Dictionary containing agents data with conversation IDs, response IDs, and query information
    """
    example_queries = [
        "Plan a 3-day trip to Paris from December 15-18, 2025. Budget is $2000. Need hotel near Eiffel Tower, round-trip flights from New York JFK, and recommend 2-3 activities per day.",
        "Find a budget hotel in Tokyo for January 5-10, 2026 under $150/night near Shibuya station, book activities including a sushi making class",
        "Search for round-trip flights from Los Angeles to London departing March 20, 2026, returning March 27, 2026. Economy class, 2 passengers. Recommend tourist attractions and museums.",
    ]

    query = example_queries[0]
    print(f"Query: {query}\n")

    result = await run_workflow_with_response_tracking(query, model=model)

    # Create output data structure
    output_data = {"agents": {}, "query": result["query"], "output": result.get("output", "")}

    # Create agent-specific mappings - now with lists of IDs
    all_agents = set(result["conversation_ids"].keys()) | set(result["response_ids"].keys())
    for agent_name in all_agents:
        output_data["agents"][agent_name] = {
            "conversation_ids": result["conversation_ids"].get(agent_name, []),
            "response_ids": result["response_ids"].get(agent_name, []),
            "response_count": len(result["response_ids"].get(agent_name, [])),
        }

    print(f"\nTotal agents tracked: {len(output_data['agents'])}")

    # Print summary of multiple responses
    print("\n=== Multi-Response Summary ===")
    for agent_name, agent_data in output_data["agents"].items():
        response_count = agent_data["response_count"]
        print(f"{agent_name}: {response_count} response(s)")

    return output_data


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