// Copyright (c) Microsoft. All rights reserved.

using Microsoft.Azure.Functions.Worker;
using Microsoft.Azure.Functions.Worker.Context.Features;
using Microsoft.Azure.Functions.Worker.Extensions.Mcp;
using Microsoft.Azure.Functions.Worker.Http;
using Microsoft.Azure.Functions.Worker.Invocation;
using Microsoft.DurableTask.Client;

namespace Microsoft.Agents.AI.Hosting.AzureFunctions;

/// <summary>
/// This implementation of function executor handles invocations using the built-in static methods for agent HTTP and entity functions.
/// </summary>
/// <remarks>By default, the Azure Functions worker generates function executor and that executor is used for function invocations.
/// But for the dummy HTTP function we create for agents (by augmenting the metadata), that executor will not have the code to handle that function since the entrypoint is a built-in static method.
/// </remarks>
internal sealed class BuiltInFunctionExecutor : IFunctionExecutor
{
    public async ValueTask ExecuteAsync(FunctionContext context)
    {
        ArgumentNullException.ThrowIfNull(context);

        // Orchestration triggers use a different input binding mechanism than other triggers.
        // The encoded orchestrator state is retrieved via BindInputAsync on the orchestration trigger binding,
        // not through IFunctionInputBindingFeature. Handle this case first to avoid unnecessary binding work.
        if (context.FunctionDefinition.EntryPoint == BuiltInFunctions.RunWorkflowOrchestrationFunctionEntryPoint)
        {
            await ExecuteOrchestrationAsync(context);
            return;
        }

        // Acquire the input binding feature (fail fast if missing rather than null-forgiving operator).
        IFunctionInputBindingFeature? functionInputBindingFeature = context.Features.Get<IFunctionInputBindingFeature>() ??
            throw new InvalidOperationException("Function input binding feature is not available on the current context.");

        FunctionInputBindingResult? inputBindingResults = await functionInputBindingFeature.BindFunctionInputAsync(context);
        if (inputBindingResults is not { Values: { } values })
        {
            throw new InvalidOperationException($"Function input binding failed for the invocation {context.InvocationId}");
        }

        HttpRequestData? httpRequestData = null;
        string? encodedEntityRequest = null;
        DurableTaskClient? durableTaskClient = null;
        ToolInvocationContext? mcpToolInvocationContext = null;

        foreach (var binding in values)
        {
            switch (binding)
            {
                case HttpRequestData request:
                    httpRequestData = request;
                    break;
                case string entityRequest:
                    encodedEntityRequest = entityRequest;
                    break;
                case DurableTaskClient client:
                    durableTaskClient = client;
                    break;
                case ToolInvocationContext toolContext:
                    mcpToolInvocationContext = toolContext;
                    break;
            }
        }

        if (durableTaskClient is null)
        {
            // This is not expected to happen since all built-in functions (other than orchestration triggers)
            // are expected to have a Durable Task client binding.
            throw new InvalidOperationException($"Durable Task client binding is missing for the invocation {context.InvocationId}.");
        }

        if (context.FunctionDefinition.EntryPoint == BuiltInFunctions.RunWorkflowOrchestrationHttpFunctionEntryPoint)
        {
            if (httpRequestData == null)
            {
                throw new InvalidOperationException($"HTTP request data binding is missing for the invocation {context.InvocationId}.");
            }

            context.GetInvocationResult().Value = await BuiltInFunctions.RunWorkflowOrchestrationHttpTriggerAsync(
                httpRequestData,
                durableTaskClient,
                context);
            return;
        }

        if (context.FunctionDefinition.EntryPoint == BuiltInFunctions.GetWorkflowStatusHttpFunctionEntryPoint)
        {
            if (httpRequestData == null)
            {
                throw new InvalidOperationException($"HTTP request data binding is missing for the invocation {context.InvocationId}.");
            }

            context.GetInvocationResult().Value = await BuiltInFunctions.GetWorkflowStatusAsync(
                httpRequestData,
                durableTaskClient,
                context);
            return;
        }

        if (context.FunctionDefinition.EntryPoint == BuiltInFunctions.RespondToWorkflowHttpFunctionEntryPoint)
        {
            if (httpRequestData == null)
            {
                throw new InvalidOperationException($"HTTP request data binding is missing for the invocation {context.InvocationId}.");
            }

            context.GetInvocationResult().Value = await BuiltInFunctions.RespondToWorkflowAsync(
                httpRequestData,
                durableTaskClient,
                context);
            return;
        }

        if (context.FunctionDefinition.EntryPoint == BuiltInFunctions.InvokeWorkflowActivityFunctionEntryPoint)
        {
            if (encodedEntityRequest is null)
            {
                throw new InvalidOperationException($"Activity trigger input binding is missing for the invocation {context.InvocationId}.");
            }

            context.GetInvocationResult().Value = await BuiltInFunctions.InvokeWorkflowActivityAsync(
                encodedEntityRequest,
                durableTaskClient,
                context);
            return;
        }

        if (context.FunctionDefinition.EntryPoint == BuiltInFunctions.RunAgentHttpFunctionEntryPoint)
        {
            if (httpRequestData == null)
            {
                throw new InvalidOperationException($"HTTP request data binding is missing for the invocation {context.InvocationId}.");
            }

            context.GetInvocationResult().Value = await BuiltInFunctions.RunAgentHttpAsync(
                httpRequestData,
                durableTaskClient,
                context);
            return;
        }

        if (context.FunctionDefinition.EntryPoint == BuiltInFunctions.RunAgentEntityFunctionEntryPoint)
        {
            if (encodedEntityRequest is null)
            {
                throw new InvalidOperationException($"Task entity dispatcher binding is missing for the invocation {context.InvocationId}.");
            }

            context.GetInvocationResult().Value = await BuiltInFunctions.InvokeAgentAsync(
                durableTaskClient,
                encodedEntityRequest,
                context);
            return;
        }

        if (context.FunctionDefinition.EntryPoint == BuiltInFunctions.RunAgentMcpToolFunctionEntryPoint)
        {
            if (mcpToolInvocationContext is null)
            {
                throw new InvalidOperationException($"MCP tool invocation context binding is missing for the invocation {context.InvocationId}.");
            }

            context.GetInvocationResult().Value =
                await BuiltInFunctions.RunMcpToolAsync(mcpToolInvocationContext, durableTaskClient, context);
            return;
        }

        if (context.FunctionDefinition.EntryPoint == BuiltInFunctions.RunWorkflowMcpToolFunctionEntryPoint)
        {
            if (mcpToolInvocationContext is null)
            {
                throw new InvalidOperationException($"MCP tool invocation context binding is missing for the invocation {context.InvocationId}.");
            }

            context.GetInvocationResult().Value = await BuiltInFunctions.RunWorkflowMcpToolAsync(
                mcpToolInvocationContext,
                durableTaskClient,
                context);
            return;
        }

        throw new InvalidOperationException($"Unsupported function entry point '{context.FunctionDefinition.EntryPoint}' for invocation {context.InvocationId}.");
    }

    private static async ValueTask ExecuteOrchestrationAsync(FunctionContext context)
    {
        BindingMetadata? orchestrationBinding = null;
        foreach (BindingMetadata binding in context.FunctionDefinition.InputBindings.Values)
        {
            if (string.Equals(binding.Type, "orchestrationTrigger", StringComparison.OrdinalIgnoreCase))
            {
                orchestrationBinding = binding;
                break;
            }
        }

        if (orchestrationBinding is null)
        {
            throw new InvalidOperationException($"Orchestration trigger binding is missing for the invocation {context.InvocationId}.");
        }

        InputBindingData<object> triggerInputData = await context.BindInputAsync<object>(orchestrationBinding);
        if (triggerInputData?.Value is not string encodedOrchestratorState)
        {
            throw new InvalidOperationException($"Orchestration history state was either missing from the input or not a string value for invocation {context.InvocationId}.");
        }

        context.GetInvocationResult().Value = BuiltInFunctions.RunWorkflowOrchestration(
            encodedOrchestratorState,
            context);
    }
}