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

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Azure.Core;
using Microsoft.Agents.AI.Workflows.Declarative.Events;
using Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests.Agents;
using Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests.Framework;
using Microsoft.Agents.AI.Workflows.Declarative.Kit;
using Microsoft.Agents.AI.Workflows.Declarative.Mcp;
using Microsoft.Extensions.AI;
using Shared.IntegrationTests;

namespace Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests;

/// <summary>
/// Integration tests for InvokeFunctionTool and InvokeMcpTool actions.
/// </summary>
public sealed class InvokeToolWorkflowTest(ITestOutputHelper output) : IntegrationTest(output)
{
    #region InvokeFunctionTool Tests

    [Theory]
    [InlineData("InvokeFunctionTool.yaml", new string[] { "GetSpecials", "GetItemPrice" }, "2.95")]
    [InlineData("InvokeFunctionToolWithApproval.yaml", new string[] { "GetItemPrice" }, "4.9")]
    public Task ValidateInvokeFunctionToolAsync(string workflowFileName, string[] expectedFunctionCalls, string? expectedResultContains) =>
        this.RunInvokeFunctionToolTestAsync(workflowFileName, expectedFunctionCalls, expectedResultContains);

    #endregion

    #region InvokeMcpTool Tests

    [Theory]
    [InlineData("InvokeMcpTool.yaml", "Azure OpenAI")]
    public Task ValidateInvokeMcpToolAsync(string workflowFileName, string? expectedResultContains) =>
        this.RunInvokeMcpToolTestAsync(workflowFileName, expectedResultContains, requireApproval: false);

    [Theory]
    [InlineData("InvokeMcpToolWithApproval.yaml", "Azure OpenAI", true)]
    [InlineData("InvokeMcpToolWithApproval.yaml", "MCP tool invocation was not approved by user", false)]
    public Task ValidateInvokeMcpToolWithApprovalAsync(string workflowFileName, string? expectedResultContains, bool approveRequest) =>
        this.RunInvokeMcpToolTestAsync(workflowFileName, expectedResultContains, requireApproval: true, approveRequest: approveRequest);

    #endregion

    #region InvokeHttpRequest Tests

    [RetryTheory(3, 5000)]
    [InlineData("HttpRequest.yaml")]
    public Task ValidateHttpRequestAsync(string workflowFileName) =>
        this.RunHttpRequestTestAsync(workflowFileName);

    #endregion

    #region InvokeFunctionTool Test Helpers

    /// <summary>
    /// Runs an InvokeFunctionTool workflow test with the specified configuration.
    /// </summary>
    private async Task RunInvokeFunctionToolTestAsync(
        string workflowFileName,
        string[] expectedFunctionCalls,
        string? expectedResultContains = null)
    {
        // Arrange
        string workflowPath = GetWorkflowPath(workflowFileName);
        IEnumerable<AIFunction> functionTools = new MenuPlugin().GetTools();
        Dictionary<string, AIFunction> functionMap = functionTools.ToDictionary(tool => tool.Name, tool => tool);
        DeclarativeWorkflowOptions workflowOptions = await this.CreateOptionsAsync(externalConversation: false);
        Workflow workflow = DeclarativeWorkflowBuilder.Build<string>(workflowPath, workflowOptions);

        WorkflowHarness harness = new(workflow, runId: Path.GetFileNameWithoutExtension(workflowPath));
        List<string> invokedFunctions = [];

        // Act - Run workflow and handle function invocations
        WorkflowEvents workflowEvents = await harness.RunWorkflowAsync("start").ConfigureAwait(false);

        while (workflowEvents.InputEvents.Count > 0)
        {
            RequestInfoEvent inputEvent = workflowEvents.InputEvents[^1];
            ExternalInputRequest? toolRequest = inputEvent.Request.Data.As<ExternalInputRequest>();
            Assert.NotNull(toolRequest);

            IList<AIContent> functionResults = await this.ProcessFunctionCallsAsync(
                toolRequest,
                functionMap,
                invokedFunctions).ConfigureAwait(false);

            ChatMessage resultMessage = new(ChatRole.Tool, functionResults);
            WorkflowEvents resumeEvents = await harness.ResumeAsync(
                inputEvent.Request.CreateResponse(new ExternalInputResponse(resultMessage))).ConfigureAwait(false);

            workflowEvents = new WorkflowEvents([.. workflowEvents.Events, .. resumeEvents.Events]);

            // Continue processing until there are no more pending input events from the resumed workflow
            if (resumeEvents.InputEvents.Count == 0)
            {
                break;
            }
        }

        // Assert - Verify function calls were made in expected order
        Assert.Equal(expectedFunctionCalls.Length, invokedFunctions.Count);
        for (int i = 0; i < expectedFunctionCalls.Length; i++)
        {
            Assert.Equal(expectedFunctionCalls[i], invokedFunctions[i]);
        }

        // Assert - Verify executor and action events
        AssertWorkflowEventsEmitted(workflowEvents);

        // Assert - Verify expected result if specified
        if (expectedResultContains is not null)
        {
            AssertResultContains(workflowEvents, expectedResultContains);
        }
    }

    /// <summary>
    /// Processes function calls from an external input request.
    /// Handles both regular function calls and approval requests.
    /// </summary>
    private async Task<IList<AIContent>> ProcessFunctionCallsAsync(
        ExternalInputRequest toolRequest,
        Dictionary<string, AIFunction> functionMap,
        List<string> invokedFunctions)
    {
        List<AIContent> results = [];

        foreach (ChatMessage message in toolRequest.AgentResponse.Messages)
        {
            // Handle approval requests if present
            foreach (ToolApprovalRequestContent approvalRequest in message.Contents.OfType<ToolApprovalRequestContent>())
            {
                this.Output.WriteLine($"APPROVAL REQUEST: {((FunctionCallContent)approvalRequest.ToolCall).Name}");
                // Auto-approve for testing
                results.Add(approvalRequest.CreateResponse(approved: true));
            }

            // Handle function calls
            foreach (FunctionCallContent functionCall in message.Contents.OfType<FunctionCallContent>())
            {
                this.Output.WriteLine($"FUNCTION CALL: {functionCall.Name}");

                if (!functionMap.TryGetValue(functionCall.Name, out AIFunction? functionTool))
                {
                    Assert.Fail($"Function not found: {functionCall.Name}");
                    continue;
                }

                invokedFunctions.Add(functionCall.Name);

                // Execute the function
                AIFunctionArguments? functionArguments = functionCall.Arguments is null
                    ? null
                    : new(functionCall.Arguments.NormalizePortableValues());

                object? result = await functionTool.InvokeAsync(functionArguments).ConfigureAwait(false);
                results.Add(new FunctionResultContent(functionCall.CallId, JsonSerializer.Serialize(result)));

                this.Output.WriteLine($"FUNCTION RESULT: {JsonSerializer.Serialize(result)}");
            }
        }

        return results;
    }

    #endregion

    #region InvokeMcpTool Test Helpers

    /// <summary>
    /// Runs an InvokeMcpTool workflow test with the specified configuration.
    /// </summary>
    private async Task RunInvokeMcpToolTestAsync(
        string workflowFileName,
        string? expectedResultContains = null,
        bool requireApproval = false,
        bool approveRequest = true)
    {
        // Arrange
        string workflowPath = GetWorkflowPath(workflowFileName);
        DefaultMcpToolHandler mcpToolProvider = new();
        DeclarativeWorkflowOptions workflowOptions = await this.CreateOptionsAsync(
            externalConversation: false,
            mcpToolProvider: mcpToolProvider);

        Workflow workflow = DeclarativeWorkflowBuilder.Build<string>(workflowPath, workflowOptions);
        WorkflowHarness harness = new(workflow, runId: Path.GetFileNameWithoutExtension(workflowPath));

        // Act - Run workflow and handle MCP tool invocations
        WorkflowEvents workflowEvents = await harness.RunWorkflowAsync("start").ConfigureAwait(false);

        while (workflowEvents.InputEvents.Count > 0)
        {
            RequestInfoEvent inputEvent = workflowEvents.InputEvents[^1];
            ExternalInputRequest? toolRequest = inputEvent.Request.Data.As<ExternalInputRequest>();
            Assert.NotNull(toolRequest);

            IList<AIContent> mcpResults = this.ProcessMcpToolRequests(
                toolRequest,
                approveRequest);

            ChatMessage resultMessage = new(ChatRole.Tool, mcpResults);
            WorkflowEvents resumeEvents = await harness.ResumeAsync(
                inputEvent.Request.CreateResponse(new ExternalInputResponse(resultMessage))).ConfigureAwait(false);

            workflowEvents = new WorkflowEvents([.. workflowEvents.Events, .. resumeEvents.Events]);

            // Continue processing until there are no more pending input events from the resumed workflow
            if (resumeEvents.InputEvents.Count == 0)
            {
                break;
            }
        }

        // Assert - Verify executor and action events
        AssertWorkflowEventsEmitted(workflowEvents);

        // Assert - Verify expected result if specified
        if (expectedResultContains is not null)
        {
            AssertResultContains(workflowEvents, expectedResultContains);
        }

        // Cleanup
        await mcpToolProvider.DisposeAsync().ConfigureAwait(false);
    }

    /// <summary>
    /// Processes MCP tool requests from an external input request.
    /// Handles approval requests for MCP tools.
    /// </summary>
    private List<AIContent> ProcessMcpToolRequests(
        ExternalInputRequest toolRequest,
        bool approveRequest)
    {
        List<AIContent> results = [];

        foreach (ChatMessage message in toolRequest.AgentResponse.Messages)
        {
            // Handle MCP approval requests if present
            foreach (ToolApprovalRequestContent approvalRequest in message.Contents.OfType<ToolApprovalRequestContent>())
            {
                this.Output.WriteLine($"MCP APPROVAL REQUEST: {approvalRequest.RequestId}");

                // Respond based on test configuration
                ToolApprovalResponseContent response = approvalRequest.CreateResponse(approved: approveRequest);
                results.Add(response);

                this.Output.WriteLine($"MCP APPROVAL RESPONSE: {(approveRequest ? "Approved" : "Rejected")}");
            }
        }

        return results;
    }

    #endregion

    #region InvokeHttpRequest Test Helpers

    /// <summary>
    /// The Azure ARM scope used to acquire bearer tokens for the HttpRequestAction
    /// integration test. Matches the URL configured in <c>HttpRequest.yaml</c>.
    /// </summary>
    private const string ArmScope = "https://management.azure.com/.default";

    /// <summary>
    /// The expected ARM endpoint. Only requests whose absolute URL exactly matches
    /// this scheme and host receive the authenticated <see cref="HttpClient"/>; all
    /// other URLs (including subdomain look-alikes such as
    /// <c>https://management.azure.com.evil.com</c>) fall through to the handler
    /// default and never see the bearer token.
    /// </summary>
    private static readonly Uri s_armEndpoint = new("https://management.azure.com/");

    /// <summary>
    /// Runs an HttpRequestAction workflow test with the specified configuration.
    /// </summary>
    /// <remarks>
    /// The workflow under test calls an authenticated Azure ARM endpoint. We acquire a
    /// single bearer token via the same Azure CLI credential used elsewhere in the
    /// integration test suite, attach it to a cached <see cref="HttpClient"/>, and route
    /// matching requests through that client via <see cref="DefaultHttpRequestHandler"/>'s
    /// <c>httpClientProvider</c> callback. The test owns the <see cref="HttpClient"/>'s
    /// lifetime and disposes it explicitly — <see cref="DefaultHttpRequestHandler"/> does
    /// not dispose provider-returned clients.
    /// </remarks>
    private async Task RunHttpRequestTestAsync(
        string workflowFileName)
    {
        // Arrange
        string workflowPath = GetWorkflowPath(workflowFileName);

        AccessToken accessToken =
            await TestAzureCliCredentials
                .CreateAzureCliCredential()
                .GetTokenAsync(new TokenRequestContext([ArmScope]), CancellationToken.None)
                .ConfigureAwait(false);

        using HttpClient authenticatedClient = new();
        authenticatedClient.DefaultRequestHeaders.Authorization =
            new AuthenticationHeaderValue("Bearer", accessToken.Token);

        await using DefaultHttpRequestHandler httpRequestHandler =
            new(httpClientProvider: (request, _) =>
            {
                if (Uri.TryCreate(request.Url, UriKind.Absolute, out Uri? requestUri) &&
                    string.Equals(requestUri.Scheme, s_armEndpoint.Scheme, StringComparison.OrdinalIgnoreCase) &&
                    string.Equals(requestUri.Host, s_armEndpoint.Host, StringComparison.OrdinalIgnoreCase))
                {
#pragma warning disable CA2025 // authenticatedClient outlives the handler (LIFO using disposal) and the workflow awaits all dispatches.
                    return Task.FromResult<HttpClient?>(authenticatedClient);
#pragma warning restore CA2025
                }

                // Fall back to the handler's internal client for any non-ARM URLs.
                return Task.FromResult<HttpClient?>(null);
            });

        DeclarativeWorkflowOptions workflowOptions = await this.CreateOptionsAsync(
            externalConversation: false,
            httpRequestHandler: httpRequestHandler);

        Workflow workflow = DeclarativeWorkflowBuilder.Build<string>(workflowPath, workflowOptions);
        WorkflowHarness harness = new(workflow, runId: Path.GetFileNameWithoutExtension(workflowPath));

        // Act
        WorkflowEvents workflowEvents = await harness.RunWorkflowAsync("start").ConfigureAwait(false);

        // Assert - Verify executor and action events
        AssertWorkflowEventsEmitted(workflowEvents);

        MessageActivityEvent? messageEvent = workflowEvents.Events
            .OfType<MessageActivityEvent>()
            .LastOrDefault();

        Assert.NotNull(messageEvent);
        Assert.NotNull(messageEvent.Message);
        Assert.True(
            Guid.TryParse(messageEvent.Message, out Guid retrievedTenantId),
            $"Expected the SendMessage payload to be a tenant GUID, but got: '{messageEvent.Message}'");
        Assert.NotEqual(Guid.Empty, retrievedTenantId);
    }

    #endregion

    #region Shared Helpers

    private static void AssertWorkflowEventsEmitted(WorkflowEvents workflowEvents)
    {
        Assert.NotEmpty(workflowEvents.ExecutorInvokeEvents);
        Assert.NotEmpty(workflowEvents.ExecutorCompleteEvents);
        Assert.NotEmpty(workflowEvents.ActionInvokeEvents);
    }

    private static void AssertResultContains(WorkflowEvents workflowEvents, string expectedResultContains)
    {
        MessageActivityEvent? messageEvent = workflowEvents.Events
            .OfType<MessageActivityEvent>()
            .LastOrDefault();

        Assert.NotNull(messageEvent);
        Assert.Contains(expectedResultContains, messageEvent.Message, StringComparison.OrdinalIgnoreCase);
    }

    private static string GetWorkflowPath(string workflowFileName) =>
        Path.Combine(Environment.CurrentDirectory, "Workflows", workflowFileName);

    #endregion
}