# Copyright (c) Microsoft. All rights reserved.

import asyncio
import logging
from collections.abc import AsyncIterable
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeVar, cast

from azure.ai.agents.models import (
    AgentsNamedToolChoiceType,
    AgentStreamEvent,
    AsyncAgentEventHandler,
    AsyncAgentRunStream,
    BaseAsyncAgentEventHandler,
    FunctionToolDefinition,
    RequiredMcpToolCall,
    ResponseFormatJsonSchemaType,
    RunStep,
    RunStepAzureAISearchToolCall,
    RunStepBingCustomSearchToolCall,
    RunStepBingGroundingToolCall,
    RunStepCodeInterpreterToolCall,
    RunStepDeepResearchToolCall,
    RunStepDeltaChunk,
    RunStepDeltaToolCallObject,
    RunStepFileSearchToolCall,
    RunStepMcpToolCall,
    RunStepMessageCreationDetails,
    RunStepOpenAPIToolCall,
    RunStepToolCallDetails,
    RunStepType,
    SubmitToolApprovalAction,
    SubmitToolOutputsAction,
    ThreadMessage,
    ThreadRun,
    ToolApproval,
    ToolDefinition,
    TruncationObject,
)
from azure.ai.agents.models._enums import MessageRole

from semantic_kernel.agents.azure_ai.agent_content_generation import (
    THREAD_MESSAGE_ID,
    generate_azure_ai_search_content,
    generate_bing_grounding_content,
    generate_code_interpreter_content,
    generate_deep_research_content,
    generate_file_search_content,
    generate_function_call_content,
    generate_function_call_streaming_content,
    generate_function_result_content,
    generate_mcp_call_content,
    generate_mcp_content,
    generate_message_content,
    generate_openapi_content,
    generate_streaming_azure_ai_search_content,
    generate_streaming_bing_grounding_content,
    generate_streaming_code_interpreter_content,
    generate_streaming_deep_research_content,
    generate_streaming_file_search_content,
    generate_streaming_mcp_call_content,
    generate_streaming_mcp_content,
    generate_streaming_message_content,
    generate_streaming_openapi_content,
    get_function_call_contents,
)
from semantic_kernel.agents.azure_ai.azure_ai_agent_utils import AzureAIAgentUtils
from semantic_kernel.agents.open_ai.assistant_content_generation import merge_streaming_function_results
from semantic_kernel.agents.open_ai.function_action_result import FunctionActionResult
from semantic_kernel.agents.open_ai.run_polling_options import RunPollingOptions
from semantic_kernel.connectors.ai.function_calling_utils import kernel_function_metadata_to_function_call_format
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
from semantic_kernel.connectors.ai.function_choice_type import FunctionChoiceType
from semantic_kernel.contents.chat_message_content import ChatMessageContent
from semantic_kernel.contents.function_call_content import FunctionCallContent
from semantic_kernel.contents.utils.author_role import AuthorRole
from semantic_kernel.exceptions.agent_exceptions import AgentInvokeException, AgentThreadOperationException
from semantic_kernel.functions import KernelArguments
from semantic_kernel.functions.kernel_function_metadata import KernelFunctionMetadata
from semantic_kernel.utils.feature_stage_decorator import experimental

if TYPE_CHECKING:
    from azure.ai.projects.aio import AIProjectClient

    from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
    from semantic_kernel.contents.chat_history import ChatHistory
    from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
    from semantic_kernel.filters.auto_function_invocation.auto_function_invocation_context import (
        AutoFunctionInvocationContext,
    )
    from semantic_kernel.kernel import Kernel

_T = TypeVar("_T", bound="AgentThreadActions")

logger: logging.Logger = logging.getLogger(__name__)


@experimental
class AgentThreadActions:
    """AzureAI Agent Thread Actions."""

    polling_status: ClassVar[list[str]] = ["queued", "in_progress", "cancelling"]
    error_message_states: ClassVar[list[str]] = ["failed", "cancelled", "expired", "incomplete"]

    # region Invocation Methods

    @classmethod
    async def invoke(
        cls: type[_T],
        *,
        agent: "AzureAIAgent",
        thread_id: str,
        arguments: KernelArguments | None = None,
        kernel: "Kernel | None" = None,
        # Run-level parameters:
        model: str | None = None,
        instructions_override: str | None = None,
        additional_instructions: str | None = None,
        additional_messages: "list[ChatMessageContent] | None" = None,
        tools: list[ToolDefinition] | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
        max_prompt_tokens: int | None = None,
        max_completion_tokens: int | None = None,
        truncation_strategy: TruncationObject | None = None,
        response_format: ResponseFormatJsonSchemaType | None = None,
        parallel_tool_calls: bool | None = None,
        metadata: dict[str, str] | None = None,
        polling_options: RunPollingOptions | None = None,
        function_choice_behavior: FunctionChoiceBehavior | None = None,
        **kwargs: Any,
    ) -> AsyncIterable[tuple[bool, "ChatMessageContent"]]:
        """Invoke the message in the thread.

        Args:
            agent: The agent to invoke.
            thread_id: The thread id.
            arguments: The kernel arguments.
            kernel: The kernel.
            model: The model.
            instructions_override: The instructions override.
            additional_instructions: The additional instructions.
            additional_messages: The additional messages to add to the thread. Only supports messages with
                role = User or Assistant.
                https://platform.openai.com/docs/api-reference/runs/createRun#runs-createrun-additional_messages
            tools: The SDK-level tools (e.g. CodeInterpreter, FileSearch, AzureAISearch). When provided,
                overrides the tools from the agent definition. Does not affect kernel function availability;
                use function_choice_behavior for that.
            temperature: The temperature.
            top_p: The top p.
            max_prompt_tokens: The max prompt tokens.
            max_completion_tokens: The max completion tokens.
            truncation_strategy: The truncation strategy.
            response_format: The response format.
            parallel_tool_calls: The parallel tool calls.
            metadata: The metadata.
            polling_options: The polling options defined at the run-level. These will override the agent-level
                polling options.
            function_choice_behavior: Controls which kernel functions are allowed to execute during this run.
                Use FunctionChoiceBehavior.Auto(filters={"included_functions": [...]}) to restrict to specific
                functions. Only Auto is supported; other types will raise an error.
            kwargs: Additional keyword arguments.

        Returns:
            A tuple of the visibility flag and the invoked message.
        """
        arguments = KernelArguments() if arguments is None else KernelArguments(**arguments, **kwargs)
        kernel = kernel or agent.kernel

        cls._validate_function_choice_behavior(function_choice_behavior)

        tools = cls._get_tools(
            agent=agent, kernel=kernel, tools_override=tools, function_choice_behavior=function_choice_behavior
        )  # type: ignore

        base_instructions = await agent.format_instructions(kernel=kernel, arguments=arguments)

        merged_instructions: str = ""
        if instructions_override is not None:
            merged_instructions = instructions_override
        elif base_instructions and additional_instructions:
            merged_instructions = f"{base_instructions}\n\n{additional_instructions}"
        else:
            merged_instructions = base_instructions or additional_instructions or ""

        run_options = cls._generate_options(
            agent=agent,
            model=model,
            additional_messages=additional_messages,
            max_completion_tokens=max_completion_tokens,
            max_prompt_tokens=max_prompt_tokens,
            temperature=temperature,
            top_p=top_p,
            metadata=metadata,
            truncation_strategy=truncation_strategy,
            response_format=response_format,
            parallel_tool_calls=parallel_tool_calls,
        )
        # Remove keys with None values.
        run_options = {k: v for k, v in run_options.items() if v is not None}

        run: ThreadRun = await agent.client.agents.runs.create(
            agent_id=agent.id,
            thread_id=thread_id,
            instructions=merged_instructions or agent.instructions,
            tools=tools,
            **run_options,
        )

        processed_step_ids = set()
        function_steps: dict[str, "FunctionCallContent"] = {}

        while run.status != "completed":
            run = await cls._poll_run_status(
                agent=agent, run=run, thread_id=thread_id, polling_options=polling_options or agent.polling_options
            )

            if run.status in cls.error_message_states:
                error_message = "None"
                if run.last_error and run.last_error.message:
                    error_message = run.last_error.message
                incomplete_details_reason = "None"
                if run.incomplete_details and run.incomplete_details.reason:
                    incomplete_details_reason = run.incomplete_details.reason
                raise AgentInvokeException(
                    f"Run failed with status: `{run.status}` for agent `{agent.name}` and thread `{thread_id}` "
                    f"with error: {error_message} and incomplete details reason: {incomplete_details_reason}"
                )

            # Check if function calling is required
            if run.status == "requires_action":
                if isinstance(run.required_action, SubmitToolOutputsAction):
                    logger.debug(
                        f"Run [{run.id}] requires tool action for agent `{agent.name}` and thread `{thread_id}`"
                    )
                    fccs = get_function_call_contents(run, function_steps)
                    if fccs:
                        logger.debug(
                            f"Yielding generate_function_call_content for agent `{agent.name}` and "
                            f"thread `{thread_id}`, visibility False"
                        )
                        yield False, generate_function_call_content(agent_name=agent.name, fccs=fccs)

                        from semantic_kernel.contents.chat_history import ChatHistory

                        chat_history = ChatHistory() if kwargs.get("chat_history") is None else kwargs["chat_history"]
                        _ = await cls._invoke_function_calls(
                            kernel=kernel,
                            fccs=fccs,
                            chat_history=chat_history,
                            arguments=arguments,
                            function_choice_behavior=function_choice_behavior,
                        )

                        tool_outputs = cls._format_tool_outputs(fccs, chat_history)
                        await agent.client.agents.runs.submit_tool_outputs(
                            run_id=run.id,
                            thread_id=thread_id,
                            tool_outputs=tool_outputs,  # type: ignore
                        )
                        logger.debug(f"Submitted tool outputs for agent `{agent.name}` and thread `{thread_id}`")
                        continue

                # Check if MCP tool approval is required
                elif isinstance(run.required_action, SubmitToolApprovalAction):
                    logger.debug(
                        f"Run [{run.id}] requires MCP tool approval for agent `{agent.name}` and thread `{thread_id}`"
                    )
                    tool_calls = run.required_action.submit_tool_approval.tool_calls
                    if not tool_calls:
                        logger.warning(f"No tool calls provided for MCP approval - cancelling run [{run.id}]")
                        await agent.client.agents.runs.cancel(run_id=run.id, thread_id=thread_id)
                        continue

                    mcp_tool_calls = [tc for tc in tool_calls if isinstance(tc, RequiredMcpToolCall)]
                    if mcp_tool_calls:
                        logger.debug(
                            f"Yielding generate_mcp_call_content for agent `{agent.name}` and "
                            f"thread `{thread_id}`, visibility False"
                        )
                        yield False, generate_mcp_call_content(agent_name=agent.name, mcp_tool_calls=mcp_tool_calls)

                        # Create tool approvals for MCP calls
                        tool_approvals = []
                        for mcp_call in mcp_tool_calls:
                            tool_approvals.append(
                                ToolApproval(
                                    tool_call_id=mcp_call.id,
                                    # TODO(evmattso): we don't support manual tool calling yet
                                    # so we always approve
                                    approve=True,
                                )
                            )

                        await agent.client.agents.runs.submit_tool_outputs(
                            run_id=run.id,
                            thread_id=thread_id,
                            tool_approvals=tool_approvals,  # type: ignore
                        )
                        logger.debug(f"Submitted MCP tool approvals for agent `{agent.name}` and thread `{thread_id}`")
                        continue

            steps: list[RunStep] = []
            async for steps_response in agent.client.agents.run_steps.list(thread_id=thread_id, run_id=run.id):
                steps.append(steps_response)
            logger.debug(f"Call for steps_response for run [{run.id}] agent `{agent.name}` and thread `{thread_id}`")

            def sort_key(step: RunStep):
                # Put tool_calls first, then message_creation.
                # If multiple steps share a type, break ties by completed_at.
                return (0 if step.type == "tool_calls" else 1, step.completed_at)

            completed_steps_to_process = sorted(
                [s for s in steps if s.completed_at is not None and s.id not in processed_step_ids],
                key=sort_key,
            )

            logger.debug(
                f"Completed steps to process for run [{run.id}] agent `{agent.name}` and thread `{thread_id}` "
                f"with length `{len(completed_steps_to_process)}`"
            )

            message_count = 0
            for completed_step in completed_steps_to_process:
                match completed_step.type:
                    case RunStepType.TOOL_CALLS:
                        logger.debug(
                            f"Entering step type tool_calls for run [{run.id}], agent `{agent.name}` and "
                            f"thread `{thread_id}`"
                        )
                        tool_call_details: RunStepToolCallDetails = cast(
                            RunStepToolCallDetails, completed_step.step_details
                        )
                        for tool_call in tool_call_details.tool_calls:
                            is_visible = False
                            content: "ChatMessageContent | None" = None
                            match tool_call.type:
                                case AgentsNamedToolChoiceType.CODE_INTERPRETER:
                                    logger.debug(
                                        f"Entering tool_calls (code_interpreter) for run [{run.id}], agent "
                                        f"`{agent.name}` and thread `{thread_id}`"
                                    )
                                    code_call: RunStepCodeInterpreterToolCall = cast(
                                        RunStepCodeInterpreterToolCall, tool_call
                                    )
                                    content = generate_code_interpreter_content(
                                        agent.name,
                                        code_call.code_interpreter.input,
                                    )
                                    is_visible = True
                                case AgentsNamedToolChoiceType.FUNCTION:
                                    logger.debug(
                                        f"Entering tool_calls (function) for run [{run.id}], agent `{agent.name}` "
                                        f"and thread `{thread_id}`"
                                    )
                                    function_step = function_steps.get(tool_call.id)
                                    assert function_step is not None  # nosec
                                    content = generate_function_result_content(
                                        agent_name=agent.name,
                                        function_step=function_step,
                                        tool_call=tool_call,  # type: ignore
                                    )
                                case (
                                    AgentsNamedToolChoiceType.BING_GROUNDING
                                    | AgentsNamedToolChoiceType.BING_CUSTOM_SEARCH
                                ):
                                    logger.debug(
                                        f"Entering tool_calls (bing_grounding/bing_custom_search) for run [{run.id}], "
                                        f"agent `{agent.name}` and thread `{thread_id}`"
                                    )
                                    # Handle both Bing grounding and custom search tool calls
                                    bing_call: RunStepBingGroundingToolCall | RunStepBingCustomSearchToolCall = cast(
                                        RunStepBingGroundingToolCall | RunStepBingCustomSearchToolCall, tool_call
                                    )
                                    content = generate_bing_grounding_content(
                                        agent_name=agent.name, bing_tool_call=bing_call
                                    )
                                case AgentsNamedToolChoiceType.AZURE_AI_SEARCH:
                                    logger.debug(
                                        f"Entering tool_calls (azure_ai_search) for run [{run.id}], agent "
                                        f" `{agent.name}` and thread `{thread_id}`"
                                    )
                                    azure_ai_search_call: RunStepAzureAISearchToolCall = cast(
                                        RunStepAzureAISearchToolCall, tool_call
                                    )
                                    content = generate_azure_ai_search_content(
                                        agent_name=agent.name, azure_ai_search_tool_call=azure_ai_search_call
                                    )
                                case AgentsNamedToolChoiceType.FILE_SEARCH:
                                    logger.debug(
                                        f"Entering tool_calls (file_search) for run [{run.id}], agent "
                                        f" `{agent.name}` and thread `{thread_id}`"
                                    )
                                    file_search_call: RunStepFileSearchToolCall = cast(
                                        RunStepFileSearchToolCall, tool_call
                                    )
                                    content = generate_file_search_content(
                                        agent_name=agent.name, file_search_tool_call=file_search_call
                                    )
                                case "openapi":
                                    logger.debug(
                                        f"Entering tool_calls (openapi) for run [{run.id}], agent "
                                        f" `{agent.name}` and thread `{thread_id}`"
                                    )
                                    openapi_tool_call: RunStepOpenAPIToolCall = cast(RunStepOpenAPIToolCall, tool_call)
                                    content = generate_openapi_content(
                                        agent_name=agent.name,
                                        openapi_tool_call=openapi_tool_call,
                                    )
                                case AgentsNamedToolChoiceType.MCP:
                                    logger.debug(
                                        f"Entering tool_calls (mcp) for run [{run.id}], agent "
                                        f" `{agent.name}` and thread `{thread_id}`"
                                    )
                                    mcp_tool_call: RunStepMcpToolCall = cast(RunStepMcpToolCall, tool_call)
                                    content = generate_mcp_content(
                                        agent_name=agent.name,
                                        mcp_tool_call=mcp_tool_call,
                                    )
                                case AgentsNamedToolChoiceType.DEEP_RESEARCH:
                                    logger.debug(
                                        f"Entering tool_calls (deep_research) for run [{run.id}], agent "
                                        f" `{agent.name}` and thread `{thread_id}`"
                                    )
                                    deep_research_call: RunStepDeepResearchToolCall = cast(
                                        RunStepDeepResearchToolCall, tool_call
                                    )
                                    content = generate_deep_research_content(
                                        agent_name=agent.name,
                                        deep_research_tool_call=deep_research_call,
                                    )

                            if content:
                                message_count += 1
                                logger.debug(
                                    f"Yielding tool_message for run [{run.id}], agent `{agent.name}`, "
                                    f"thread `{thread_id}`, message count `{message_count}`, "
                                    f"is_visible `{is_visible}`"
                                )
                                yield is_visible, content
                    case RunStepType.MESSAGE_CREATION:
                        logger.debug(
                            f"Entering message_creation for run [{run.id}], agent `{agent.name}` and thread "
                            f"`{thread_id}`"
                        )
                        message_call_details: RunStepMessageCreationDetails = cast(
                            RunStepMessageCreationDetails, completed_step.step_details
                        )
                        message = await cls._retrieve_message(
                            agent=agent,
                            thread_id=thread_id,
                            message_id=message_call_details.message_creation.message_id,  # type: ignore
                        )
                        if message:
                            content = generate_message_content(agent.name, message, completed_step)
                            if content and len(content.items) > 0:
                                message_count += 1
                                logger.debug(
                                    f"Yielding message_creation for run [{run.id}], agent `{agent.name}`, "
                                    f"thread `{thread_id}`, message count `{message_count}`, is_visible `True`"
                                )
                                yield True, content
                processed_step_ids.add(completed_step.id)

    @classmethod
    async def invoke_stream(
        cls: type[_T],
        *,
        agent: "AzureAIAgent",
        thread_id: str,
        additional_instructions: str | None = None,
        additional_messages: "list[ChatMessageContent] | None" = None,
        arguments: KernelArguments | None = None,
        instructions_override: str | None = None,
        kernel: "Kernel | None" = None,
        metadata: dict[str, str] | None = None,
        model: str | None = None,
        max_prompt_tokens: int | None = None,
        max_completion_tokens: int | None = None,
        output_messages: list[ChatMessageContent] | None = None,
        parallel_tool_calls: bool | None = None,
        response_format: ResponseFormatJsonSchemaType | None = None,
        tools: list[ToolDefinition] | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
        truncation_strategy: TruncationObject | None = None,
        function_choice_behavior: FunctionChoiceBehavior | None = None,
        **kwargs: Any,
    ) -> AsyncIterable["StreamingChatMessageContent"]:
        """Invoke the agent stream and yield ChatMessageContent continuously.

        Args:
            agent: The agent to invoke.
            thread_id: The thread id.
            additional_instructions: The additional instructions.
            additional_messages: The additional messages to add to the thread. Only supports messages with
                role = User or Assistant.
                https://platform.openai.com/docs/api-reference/runs/createRun
            arguments: The kernel arguments.
            instructions_override: The instructions override.
            kernel: The kernel.
            metadata: The metadata.
            model: The model.
            max_prompt_tokens: The max prompt tokens.
            max_completion_tokens: The max completion tokens.
            output_messages: The output messages received from the agent. These are full content messages
                formed from the streamed chunks.
            parallel_tool_calls: Whether to configure parallel tool calls.
            response_format: The response format.
            tools: The SDK-level tools (e.g. CodeInterpreter, FileSearch, AzureAISearch). When provided,
                overrides the tools from the agent definition. Does not affect kernel function availability;
                use function_choice_behavior for that.
            temperature: The temperature.
            top_p: The top p.
            truncation_strategy: The truncation strategy.
            function_choice_behavior: Controls which kernel functions are allowed to execute during this run.
                Use FunctionChoiceBehavior.Auto(filters={"included_functions": [...]}) to restrict to specific
                functions. Only Auto is supported; other types will raise an error.
            kwargs: Additional keyword arguments.

        Returns:
            An async iterable of StreamingChatMessageContent.
        """
        arguments = KernelArguments() if arguments is None else KernelArguments(**arguments, **kwargs)
        kernel = kernel or agent.kernel
        arguments = agent._merge_arguments(arguments)

        cls._validate_function_choice_behavior(function_choice_behavior)

        tools = cls._get_tools(
            agent=agent, kernel=kernel, tools_override=tools, function_choice_behavior=function_choice_behavior
        )  # type: ignore

        base_instructions = await agent.format_instructions(kernel=kernel, arguments=arguments)

        merged_instructions: str = ""
        if instructions_override is not None:
            merged_instructions = instructions_override
        elif base_instructions and additional_instructions:
            merged_instructions = f"{base_instructions}\n\n{additional_instructions}"
        else:
            merged_instructions = base_instructions or additional_instructions or ""

        run_options = cls._generate_options(
            agent=agent,
            model=model,
            additional_messages=additional_messages,
            max_completion_tokens=max_completion_tokens,
            max_prompt_tokens=max_prompt_tokens,
            temperature=temperature,
            top_p=top_p,
            metadata=metadata,
            truncation_strategy=truncation_strategy,
            response_format=response_format,
            parallel_tool_calls=parallel_tool_calls,
        )
        run_options = {k: v for k, v in run_options.items() if v is not None}

        stream: AsyncAgentRunStream = await agent.client.agents.runs.stream(
            agent_id=agent.id,
            thread_id=thread_id,
            instructions=merged_instructions or agent.instructions,
            tools=tools,
            **run_options,
        )

        function_steps: dict[str, FunctionCallContent] = {}
        active_messages: dict[str, RunStep] = {}

        async for content in cls._process_stream_events(
            stream=stream,
            agent=agent,
            thread_id=thread_id,
            output_messages=output_messages,
            kernel=kernel,
            arguments=arguments,
            function_steps=function_steps,
            active_messages=active_messages,
            function_choice_behavior=function_choice_behavior,
        ):
            if content:
                yield content

    @classmethod
    async def _process_stream_events(
        cls: type[_T],
        stream: AsyncAgentRunStream,
        agent: "AzureAIAgent",
        thread_id: str,
        kernel: "Kernel",
        arguments: KernelArguments,
        function_steps: dict[str, FunctionCallContent],
        active_messages: dict[str, RunStep],
        output_messages: "list[ChatMessageContent] | None" = None,
        function_choice_behavior: FunctionChoiceBehavior | None = None,
    ) -> AsyncIterable["StreamingChatMessageContent"]:
        """Process events from the main stream and delegate tool output handling as needed."""
        thread_msg_id = None
        while True:
            # Use 'async with' only if the stream supports async context management (main agent stream).
            # Tool output handlers only support async iteration, not context management.
            if hasattr(stream, "__aenter__") and hasattr(stream, "__aexit__"):
                async with stream as response_stream:
                    stream_iter = response_stream
            else:
                stream_iter = stream
            async for event_type, event_data, _ in stream_iter:
                if event_type == AgentStreamEvent.THREAD_RUN_CREATED:
                    run = event_data
                    logger.info(f"Assistant run created with ID: {run.id}")

                elif event_type == AgentStreamEvent.THREAD_RUN_IN_PROGRESS:
                    run_step = cast(RunStep, event_data)
                    logger.info(f"Assistant run in progress with ID: {run_step.id}")

                elif event_type == AgentStreamEvent.THREAD_MESSAGE_CREATED:
                    # Keep the current message id stable unless a new one arrives
                    if thread_msg_id != event_data.id:
                        thread_msg_id = event_data.id
                    logger.info(f"Assistant message created with ID: {thread_msg_id}")

                elif event_type == AgentStreamEvent.THREAD_MESSAGE_DELTA:
                    yield generate_streaming_message_content(agent.name, event_data, thread_msg_id)

                elif event_type == AgentStreamEvent.THREAD_RUN_STEP_COMPLETED:
                    step_completed = cast(RunStep, event_data)
                    logger.info(f"Run step completed with ID: {step_completed.id}")
                    if isinstance(step_completed.step_details, RunStepMessageCreationDetails):
                        msg_id = step_completed.step_details.message_creation.message_id
                        active_messages.setdefault(msg_id, step_completed)

                elif event_type == AgentStreamEvent.THREAD_RUN_STEP_DELTA:
                    run_step_event: RunStepDeltaChunk = event_data
                    details = run_step_event.delta.step_details
                    if not details:
                        continue
                    if isinstance(details, RunStepDeltaToolCallObject) and details.tool_calls:
                        content_is_visible = False
                        for tool_call in details.tool_calls:
                            logger.debug(
                                f"Generating content for tool call type `{tool_call.type}`, agent `{agent.name}` and "
                                f"thread `{thread_id}` with tool call details: {details}"
                            )
                            content = None
                            match tool_call.type:
                                # Function Calling-related content is emitted as a single message
                                # via the `on_intermediate_message` callback.
                                case AgentsNamedToolChoiceType.CODE_INTERPRETER:
                                    content = generate_streaming_code_interpreter_content(agent.name, details)
                                    content_is_visible = True
                                case (
                                    AgentsNamedToolChoiceType.BING_GROUNDING
                                    | AgentsNamedToolChoiceType.BING_CUSTOM_SEARCH
                                ):
                                    content = generate_streaming_bing_grounding_content(
                                        agent_name=agent.name, step_details=details
                                    )
                                case AgentsNamedToolChoiceType.AZURE_AI_SEARCH:
                                    content = generate_streaming_azure_ai_search_content(
                                        agent_name=agent.name, step_details=details
                                    )
                                case AgentsNamedToolChoiceType.FILE_SEARCH:
                                    content = generate_streaming_file_search_content(
                                        agent_name=agent.name, step_details=details
                                    )
                                case "openapi":
                                    # There's no enum for OpenAPI tool calls as part of `AgentsNamedToolChoiceType`
                                    # so we handle it separately.
                                    content = generate_streaming_openapi_content(
                                        agent_name=agent.name, step_details=details
                                    )
                                case AgentsNamedToolChoiceType.MCP:
                                    content = generate_streaming_mcp_content(
                                        agent_name=agent.name, step_details=details
                                    )
                                case AgentsNamedToolChoiceType.DEEP_RESEARCH:
                                    content = generate_streaming_deep_research_content(
                                        agent_name=agent.name, step_details=details
                                    )
                            if content:
                                if thread_msg_id and THREAD_MESSAGE_ID not in content.metadata:
                                    content.metadata[THREAD_MESSAGE_ID] = thread_msg_id
                                if output_messages is not None:
                                    output_messages.append(content)
                                if content_is_visible:
                                    yield content

                elif event_type == AgentStreamEvent.THREAD_RUN_REQUIRES_ACTION:
                    logger.debug(
                        f"Entering step type {event_type}, agent `{agent.name}` and "
                        f"thread `{thread_id}` with event data: {event_data}"
                    )
                    run = cast(ThreadRun, event_data)

                    # Check if this is a function call request
                    if isinstance(run.required_action, SubmitToolOutputsAction):
                        action_result = await cls._handle_streaming_requires_action(
                            agent_name=agent.name,
                            kernel=kernel,
                            run=run,
                            function_steps=function_steps,
                            arguments=arguments,
                            function_choice_behavior=function_choice_behavior,
                        )
                        if action_result is None:
                            raise RuntimeError(
                                f"Function call required but no function steps found for agent `{agent.name}` "
                                f"thread: {thread_id}."
                            )

                        for content in (
                            action_result.function_call_streaming_content,
                            action_result.function_result_streaming_content,
                        ):
                            if content and output_messages is not None:
                                if thread_msg_id and THREAD_MESSAGE_ID not in content.metadata:
                                    content.metadata[THREAD_MESSAGE_ID] = thread_msg_id
                                output_messages.append(content)

                        handler: BaseAsyncAgentEventHandler = AsyncAgentEventHandler()
                        await agent.client.agents.runs.submit_tool_outputs_stream(
                            run_id=run.id,
                            thread_id=thread_id,
                            tool_outputs=action_result.tool_outputs,  # type: ignore
                            event_handler=handler,
                        )
                        # Pass the handler to the stream to continue processing
                        stream = handler  # type: ignore

                        logger.debug(
                            f"Submitted tool outputs stream for agent `{agent.name}` and "
                            f"thread `{thread_id}` and run id `{run.id}`"
                        )
                        break

                    # Check if this is an MCP tool approval request
                    elif isinstance(run.required_action, SubmitToolApprovalAction):
                        tool_calls = run.required_action.submit_tool_approval.tool_calls
                        if not tool_calls:
                            logger.warning(f"No tool calls provided for MCP approval - cancelling run [{run.id}]")
                            await agent.client.agents.runs.cancel(run_id=run.id, thread_id=thread_id)
                            break

                        mcp_tool_calls = [tc for tc in tool_calls if isinstance(tc, RequiredMcpToolCall)]
                        if mcp_tool_calls:
                            logger.debug(
                                f"Processing MCP tool approvals for agent `{agent.name}` and "
                                f"thread `{thread_id}` and run id `{run.id}`"
                            )

                            if output_messages is not None:
                                content = generate_streaming_mcp_call_content(
                                    agent_name=agent.name, mcp_tool_calls=mcp_tool_calls
                                )
                                if content:
                                    if thread_msg_id and THREAD_MESSAGE_ID not in content.metadata:
                                        content.metadata[THREAD_MESSAGE_ID] = thread_msg_id
                                    output_messages.append(content)

                            # Create tool approvals for MCP calls
                            tool_approvals = []
                            for mcp_call in mcp_tool_calls:
                                tool_approvals.append(
                                    ToolApproval(
                                        tool_call_id=mcp_call.id,
                                        approve=True,
                                        # Note: headers would need to be provided by the MCP tool configuration
                                        # This is a simplified implementation
                                        headers={},
                                    )
                                )

                            handler: BaseAsyncAgentEventHandler = AsyncAgentEventHandler()  # type: ignore
                            await agent.client.agents.runs.submit_tool_outputs_stream(
                                run_id=run.id,
                                thread_id=thread_id,
                                tool_approvals=tool_approvals,  # type: ignore
                                event_handler=handler,
                            )
                            # Pass the handler to the stream to continue processing
                            stream = handler  # type: ignore

                            logger.debug(
                                f"Submitted MCP tool approvals stream for agent `{agent.name}` and "
                                f"thread `{thread_id}` and run id `{run.id}`"
                            )
                            break

                elif event_type == AgentStreamEvent.THREAD_RUN_COMPLETED:
                    logger.debug(
                        f"Entering step type {event_type}, agent `{agent.name}` and "
                        f"thread `{thread_id}` and run id `{run.id}`"
                    )
                    run = cast(ThreadRun, event_data)
                    logger.info(f"Run completed with ID: {run.id}")
                    if active_messages:
                        for msg_id, step in active_messages.items():
                            message = await cls._retrieve_message(agent=agent, thread_id=thread_id, message_id=msg_id)
                            if message and hasattr(message, "content"):
                                final_content = generate_message_content(agent.name, message, step)
                                if output_messages is not None:
                                    output_messages.append(final_content)
                    return

                elif event_type == AgentStreamEvent.THREAD_RUN_FAILED:
                    run_failed = cast(ThreadRun, event_data)
                    error_message = "None"
                    if run_failed.last_error and run_failed.last_error.message:
                        error_message = run_failed.last_error.message
                    incomplete_details_reason = "None"
                    if run_failed.incomplete_details and run_failed.incomplete_details.reason:
                        incomplete_details_reason = run_failed.incomplete_details.reason
                    raise RuntimeError(
                        f"Run failed with status: `{run_failed.status}` for agent `{agent.name}` "
                        f"thread `{thread_id}` with error: {error_message} and incomplete details reason: "
                        f"{incomplete_details_reason}"
                    )
            else:
                break
        return

    # endregion

    # region Messaging Handling Methods

    @classmethod
    async def create_thread(
        cls: type[_T],
        client: "AIProjectClient",
        **kwargs: Any,
    ) -> str:
        """Create a thread.

        Args:
            client: The client to use to create the thread.
            kwargs: Additional keyword arguments.

        Returns:
            The ID of the created thread.
        """
        thread = await client.agents.threads.create(**kwargs)
        return thread.id

    @classmethod
    async def create_message(
        cls: type[_T],
        client: "AIProjectClient",
        thread_id: str,
        message: "str | ChatMessageContent",
        **kwargs: Any,
    ) -> "ThreadMessage | None":
        """Create a message in the thread.

        Args:
            client: The client to use to create the message.
            thread_id: The ID of the thread to create the message in.
            message: The message to create.
            kwargs: Additional keyword arguments.

        Returns:
            The created message.
        """
        if isinstance(message, str):
            message = ChatMessageContent(role=AuthorRole.USER, content=message)

        if any(isinstance(item, FunctionCallContent) for item in message.items):
            return None

        if not message.content.strip():
            return None

        return await client.agents.messages.create(
            thread_id=thread_id,
            role=MessageRole.USER if message.role == AuthorRole.USER else MessageRole.AGENT,
            content=message.content,
            attachments=AzureAIAgentUtils.get_attachments(message),
            metadata=AzureAIAgentUtils.get_metadata(message),
            **kwargs,
        )

    @classmethod
    async def get_messages(
        cls: type[_T],
        client: "AIProjectClient",
        thread_id: str,
        sort_order: Literal["asc", "desc"] = "desc",
    ) -> AsyncIterable["ChatMessageContent"]:
        """Get messages from a thread.

        Args:
            client: The client to use to get the messages.
            thread_id: The ID of the thread to get the messages from.
            sort_order: The order to sort the messages in.

        Yields:
            An AsyncIterable of ChatMessageContent that includes the thread messages.

        Raises:
            AgentThreadOperationException: If the messages cannot be retrieved.
        """
        try:
            async for message in client.agents.messages.list(
                thread_id=thread_id,
                run_id=None,
                limit=None,
                order=sort_order,
                before=None,
            ):
                agent_id = (message.agent_id or message.metadata.get("agent_id") or "").strip() or "agent"
                yield generate_message_content(agent_id, message)
        except Exception as e:
            logger.error(f"Failed to retrieve messages for thread {thread_id}: {e}")
            raise AgentThreadOperationException(f"Failed to retrieve messages for thread `{thread_id}`.") from e

    # endregion

    # region Internal Methods

    @classmethod
    def _merge_options(
        cls: type[_T],
        *,
        agent: "AzureAIAgent",
        model: str | None = None,
        response_format: ResponseFormatJsonSchemaType | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
        metadata: dict[str, str] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Merge run-time options with the agent-level options.

        Run-level parameters take precedence.
        """
        return {
            "model": model if model is not None else agent.definition.model,
            "response_format": response_format if response_format is not None else agent.definition.response_format,
            "temperature": temperature if temperature is not None else None,
            "top_p": top_p if top_p is not None else None,
            "metadata": metadata if metadata is not None else agent.definition.metadata,
            **kwargs,
        }

    @classmethod
    def _generate_options(cls: type[_T], **kwargs: Any) -> dict[str, Any]:
        """Generate a dictionary of options that can be passed directly to create_run."""
        merged = cls._merge_options(**kwargs)
        truncation_strategy = merged.get("truncation_strategy", None)
        max_completion_tokens = merged.get("max_completion_tokens", None)
        max_prompt_tokens = merged.get("max_prompt_tokens", None)
        parallel_tool_calls = merged.get("parallel_tool_calls_enabled", None)
        additional_messages = cls._translate_additional_messages(merged.get("additional_messages", None))
        return {
            "model": merged.get("model"),
            "top_p": merged.get("top_p"),
            "response_format": merged.get("response_format"),
            "temperature": merged.get("temperature"),
            "truncation_strategy": truncation_strategy,
            "metadata": merged.get("metadata"),
            "max_completion_tokens": max_completion_tokens,
            "max_prompt_tokens": max_prompt_tokens,
            "parallel_tool_calls": parallel_tool_calls,
            "additional_messages": additional_messages,
        }

    @classmethod
    def _translate_additional_messages(
        cls: type[_T], messages: "list[ChatMessageContent] | None"
    ) -> list[ThreadMessage] | None:
        """Translate additional messages to the required format."""
        if not messages:
            return None
        return AzureAIAgentUtils.get_thread_messages(messages)

    @classmethod
    def _prepare_tool_definition(cls: type[_T], tool: dict | ToolDefinition) -> dict | ToolDefinition:
        """Prepare the tool definition."""
        if tool.get("type") == "openapi" and "openapi" in tool:
            openapi_data = dict(tool["openapi"])
            openapi_data.pop("functions", None)
            tool = dict(tool)
            tool["openapi"] = openapi_data
        return tool

    @staticmethod
    def _deduplicate_tools(existing_tools: list[dict], new_tools: list[dict]) -> list[dict]:
        existing_names = {
            tool["function"]["name"] for tool in existing_tools if "function" in tool and "name" in tool["function"]
        }
        return [tool for tool in new_tools if tool.get("function", {}).get("name") not in existing_names]

    @staticmethod
    def _validate_function_choice_behavior(
        function_choice_behavior: FunctionChoiceBehavior | None,
    ) -> None:
        """Validate the function choice behavior is compatible with agent invocations."""
        if function_choice_behavior is None:
            return
        if function_choice_behavior.type_ != FunctionChoiceType.AUTO:
            raise AgentInvokeException(
                f"FunctionChoiceBehavior with type '{function_choice_behavior.type_}' is not supported for agent "
                "invocations. Use FunctionChoiceBehavior.Auto(filters=...) to control which kernel functions "
                "are available."
            )
        if not function_choice_behavior.auto_invoke_kernel_functions:
            raise AgentInvokeException(
                "FunctionChoiceBehavior.Auto(auto_invoke=False) is not supported for agent invocations. "
                "The agent run loop manages tool invocation; disabling auto_invoke is not compatible."
            )
        valid_filter_keys: set[str] = {
            "excluded_plugins",
            "included_plugins",
            "excluded_functions",
            "included_functions",
        }
        if function_choice_behavior.filters is not None:
            if not function_choice_behavior.filters:
                raise AgentInvokeException(
                    "FunctionChoiceBehavior filters must not be empty. Provide at least one filter key "
                    f"from {sorted(valid_filter_keys)}, or omit filters entirely to include all "
                    "kernel functions."
                )
            unknown_keys = {str(k) for k in function_choice_behavior.filters} - valid_filter_keys
            if unknown_keys:
                raise AgentInvokeException(
                    f"Unknown filter key(s): {sorted(unknown_keys)}. "
                    f"Valid filter keys are: {sorted(valid_filter_keys)}."
                )

    @classmethod
    def _get_tools(
        cls: type[_T],
        agent: "AzureAIAgent",
        kernel: "Kernel",
        tools_override: list[ToolDefinition] | None = None,
        function_choice_behavior: FunctionChoiceBehavior | None = None,
    ) -> list[dict[str, Any] | ToolDefinition]:
        """Get the tools for the agent.

        Args:
            agent: The agent instance.
            kernel: The kernel to use for function metadata.
            tools_override: When provided, overrides agent.definition.tools (SDK-level tools only).
            function_choice_behavior: When provided, filters which kernel functions are included.
        """
        tools: list[Any] = list(tools_override) if tools_override is not None else list(agent.definition.tools)

        # Always validate against the full kernel function list to catch truly
        # unregistered functions, regardless of FCB filtering.
        all_funcs = kernel.get_full_list_of_function_metadata()
        cls._validate_function_tools_registered(tools, all_funcs)

        # Determine which kernel functions to advertise based on function_choice_behavior
        if function_choice_behavior is not None and not function_choice_behavior.enable_kernel_functions:
            funcs: list[KernelFunctionMetadata] = []
        elif function_choice_behavior is not None and function_choice_behavior.filters:
            funcs = kernel.get_list_of_function_metadata(function_choice_behavior.filters)
        else:
            funcs = all_funcs
        dict_defs = [kernel_function_metadata_to_function_call_format(f) for f in funcs]
        deduped_defs = cls._deduplicate_tools(tools, dict_defs)
        tools.extend(deduped_defs)
        return [cls._prepare_tool_definition(tool) for tool in tools]

    @staticmethod
    def _validate_function_tools_registered(
        tools: list[Any],
        funcs: list[Any],
    ) -> None:
        """Validate that all function tools are registered with the kernel."""
        function_tool_names = set()
        for tool in tools:
            if isinstance(tool, FunctionToolDefinition):
                agent_tool_func_name = getattr(tool.function, "name", None)
                if agent_tool_func_name:
                    function_tool_names.add(agent_tool_func_name)

        kernel_function_names = set()
        for f in funcs:
            kernel_func_name = (
                f.fully_qualified_name
                if isinstance(f, KernelFunctionMetadata)
                else getattr(f, "full_qualified_name", None)
            )
            if kernel_func_name:
                kernel_function_names.add(kernel_func_name)

        missing_functions = function_tool_names - kernel_function_names
        if missing_functions:
            raise AgentInvokeException(
                f"The following function tool(s) are defined on the agent but missing from the kernel: "
                f"{sorted(missing_functions)}. "
                f"Please ensure all required tools are registered with the kernel."
            )

    @classmethod
    async def _poll_run_status(
        cls: type[_T], agent: "AzureAIAgent", run: ThreadRun, thread_id: str, polling_options: RunPollingOptions
    ) -> ThreadRun:
        """Poll the run status."""
        logger.info(f"Polling run status: {run.id}, threadId: {thread_id}")
        try:
            run = await asyncio.wait_for(
                cls._poll_loop(agent=agent, run=run, thread_id=thread_id, polling_options=polling_options),
                timeout=polling_options.run_polling_timeout.total_seconds(),
            )
        except asyncio.TimeoutError:
            timeout_duration = polling_options.run_polling_timeout
            error_message = (
                f"Polling timed out for run id: `{run.id}` and thread id: `{thread_id}` "
                f"after waiting {timeout_duration}."
            )
            logger.error(error_message)
            raise AgentInvokeException(error_message)
        logger.info(f"Polled run status: {run.status}, {run.id}, threadId: {thread_id}")
        return run

    @classmethod
    async def _poll_loop(
        cls: type[_T], agent: "AzureAIAgent", run: ThreadRun, thread_id: str, polling_options: RunPollingOptions
    ) -> ThreadRun:
        """Continuously poll the run status until it is no longer pending."""
        count = 0
        while True:
            await asyncio.sleep(polling_options.get_polling_interval(count).total_seconds())
            count += 1
            try:
                run = await agent.client.agents.runs.get(run_id=run.id, thread_id=thread_id)
            except Exception as e:
                logger.warning(f"Failed to retrieve run for run id: `{run.id}` and thread id: `{thread_id}`: {e}")
            if run.status not in cls.polling_status:
                break
        return run

    @classmethod
    async def _retrieve_message(
        cls: type[_T], agent: "AzureAIAgent", thread_id: str, message_id: str
    ) -> ThreadMessage | None:
        """Retrieve a message from a thread."""
        message: ThreadMessage | None = None
        count = 0
        max_retries = 3
        while count < max_retries:
            try:
                message = await agent.client.agents.messages.get(thread_id=thread_id, message_id=message_id)
                break
            except Exception as ex:
                logger.error(f"Failed to retrieve message {message_id} from thread {thread_id}: {ex}")
                count += 1
                if count >= max_retries:
                    logger.error(
                        f"Max retries reached. Unable to retrieve message {message_id} from thread {thread_id}."
                    )
                    break
                backoff_time: float = agent.polling_options.message_synchronization_delay.total_seconds() * (2**count)
                await asyncio.sleep(backoff_time)
        return message

    @classmethod
    async def _invoke_function_calls(
        cls: type[_T],
        kernel: "Kernel",
        fccs: list["FunctionCallContent"],
        chat_history: "ChatHistory",
        arguments: KernelArguments,
        function_choice_behavior: FunctionChoiceBehavior | None = None,
    ) -> list["AutoFunctionInvocationContext | None"]:
        """Invoke the function calls."""
        return await asyncio.gather(
            *[
                kernel.invoke_function_call(
                    function_call=function_call,
                    chat_history=chat_history,
                    arguments=arguments,
                    function_behavior=function_choice_behavior,
                )
                for function_call in fccs
            ],
        )

    @classmethod
    def _format_tool_outputs(
        cls: type[_T], fccs: list["FunctionCallContent"], chat_history: "ChatHistory"
    ) -> list[dict[str, str]]:
        """Format the tool outputs for submission."""
        from semantic_kernel.contents.function_result_content import FunctionResultContent

        tool_call_lookup = {
            tool_call.id: tool_call
            for message in chat_history.messages
            for tool_call in message.items
            if isinstance(tool_call, FunctionResultContent) and tool_call.id is not None
        }
        return [
            {"tool_call_id": fcc.id, "output": str(tool_call_lookup[fcc.id].result)}
            for fcc in fccs
            if fcc.id in tool_call_lookup
        ]

    @classmethod
    async def _handle_streaming_requires_action(
        cls: type[_T],
        agent_name: str,
        kernel: "Kernel",
        run: ThreadRun,
        function_steps: dict[str, "FunctionCallContent"],
        arguments: KernelArguments,
        function_choice_behavior: FunctionChoiceBehavior | None = None,
        **kwargs: Any,
    ) -> FunctionActionResult | None:
        """Handle the requires action event for a streaming run."""
        fccs = get_function_call_contents(run, function_steps)
        if fccs:
            function_call_streaming_content = generate_function_call_streaming_content(agent_name=agent_name, fccs=fccs)
            from semantic_kernel.contents.chat_history import ChatHistory

            chat_history = ChatHistory() if kwargs.get("chat_history") is None else kwargs["chat_history"]
            results = await cls._invoke_function_calls(
                kernel=kernel,
                fccs=fccs,
                chat_history=chat_history,
                arguments=arguments,
                function_choice_behavior=function_choice_behavior,
            )

            function_result_streaming_content = merge_streaming_function_results(
                messages=chat_history.messages[-len(results) :],
                name=agent_name,
            )
            tool_outputs = cls._format_tool_outputs(fccs, chat_history)
            return FunctionActionResult(
                function_call_streaming_content,
                function_result_streaming_content,
                tool_outputs,
            )
        return None

    # endregion
