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

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Agents.AI.Workflows.Declarative.Events;
using Microsoft.Agents.AI.Workflows.Declarative.Extensions;
using Microsoft.Agents.AI.Workflows.Declarative.Interpreter;
using Microsoft.Agents.AI.Workflows.Declarative.Kit;
using Microsoft.Agents.AI.Workflows.Declarative.PowerFx;
using Microsoft.Agents.ObjectModel;
using Microsoft.Extensions.AI;
using Microsoft.Shared.Diagnostics;

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

/// <summary>
/// Executor for the <see cref="InvokeMcpTool"/> action.
/// This executor invokes MCP tools on remote servers and handles approval flows.
/// </summary>
internal sealed class InvokeMcpToolExecutor(
    InvokeMcpTool model,
    IMcpToolHandler mcpToolHandler,
    ResponseAgentProvider agentProvider,
    WorkflowFormulaState state) :
    DeclarativeActionExecutor<InvokeMcpTool>(model, state)
{
    private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshots);
    private const string LegacyApprovalSnapshotStateKey = "_approvalSnapshot";

    /// <summary>
    /// Snapshots of evaluated parameters captured at approval-request time, keyed by
    /// per-invocation request id. Each pending approval lives here until the matching
    /// response is captured.
    /// </summary>
    private readonly ConcurrentDictionary<string, ApprovalSnapshot> _approvalSnapshots = new(StringComparer.Ordinal);

    /// <summary>
    /// Step identifiers for the MCP tool invocation workflow.
    /// </summary>
    public static class Steps
    {
        /// <summary>
        /// Step for waiting for external input (approval or direct response).
        /// </summary>
        public static string ExternalInput(string id) => $"{id}_{nameof(ExternalInput)}";

        /// <summary>
        /// Step for resuming after receiving external input.
        /// </summary>
        public static string Resume(string id) => $"{id}_{nameof(Resume)}";
    }

    /// <summary>
    /// Determines if the message indicates external input is required.
    /// </summary>
    public static bool RequiresInput(object? message) =>
        message is ExternalInputRequest || (message is PortableValue pv && pv.IsType(out ExternalInputRequest? _));

    /// <summary>
    /// Determines if the message indicates no external input is required.
    /// </summary>
    public static bool RequiresNothing(object? message) =>
        message is ActionExecutorResult || (message is PortableValue pv && pv.IsType(out ActionExecutorResult? _));

    /// <inheritdoc/>
    protected override bool EmitResultEvent => false;

    /// <inheritdoc/>
    protected override bool IsDiscreteAction => false;

    /// <inheritdoc/>
    [SendsMessage(typeof(ExternalInputRequest))]
    protected override async ValueTask<object?> ExecuteAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
    {
        string serverUrl = this.GetServerUrl();
        string? serverLabel = this.GetServerLabel();
        string toolName = this.GetToolName();
        bool requireApproval = this.GetRequireApproval();
        Dictionary<string, object?>? arguments = this.GetArguments();
        Dictionary<string, string>? headers = this.GetHeaders();
        string? connectionName = this.GetConnectionName();

        if (requireApproval)
        {
            // Per-invocation request id stamped on the outbound content.
            string requestId = Guid.NewGuid().ToString("N");

            // Capture the evaluated parameters keyed by request id; the matching response
            // resumes from this snapshot.
            this._approvalSnapshots[requestId] = new ApprovalSnapshot(serverUrl, serverLabel, toolName, arguments, connectionName);

            // Create tool call content for approval request.
            // Transport headers (e.g. Authorization) are intentionally excluded from the
            // approval event: they must not cross into the externally-surfaced approval request.
            McpServerToolCallContent toolCall = new(requestId, toolName, serverLabel ?? serverUrl)
            {
                Arguments = arguments
            };

            ToolApprovalRequestContent approvalRequest = new(requestId, toolCall);

            ChatMessage requestMessage = new(ChatRole.Assistant, [approvalRequest]);
            AgentResponse agentResponse = new([requestMessage]);

            // Yield to the caller for approval
            ExternalInputRequest inputRequest = new(agentResponse);
            await context.SendMessageAsync(inputRequest, cancellationToken).ConfigureAwait(false);

            return default;
        }

        // No approval required - invoke the tool directly
        McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync(
            serverUrl,
            serverLabel,
            toolName,
            arguments,
            headers,
            connectionName,
            cancellationToken).ConfigureAwait(false);

        await this.ProcessResultAsync(context, resultContent, cancellationToken).ConfigureAwait(false);

        // Signal completion so the workflow routes via RequiresNothing
        await context.SendResultMessageAsync(this.Id, result: null, cancellationToken).ConfigureAwait(false);

        return default;
    }

    /// <summary>
    /// Captures the external input response and processes the MCP tool result.
    /// </summary>
    /// <param name="context">The workflow context.</param>
    /// <param name="response">The external input response.</param>
    /// <param name="cancellationToken">A cancellation token.</param>
    /// <returns>A <see cref="ValueTask"/> representing the asynchronous operation.</returns>
    public async ValueTask CaptureResponseAsync(
        IWorkflowContext context,
        ExternalInputResponse response,
        CancellationToken cancellationToken)
    {
        ToolApprovalResponseContent? approvalResponse = response.Messages
            .SelectMany(m => m.Contents)
            .OfType<ToolApprovalResponseContent>()
            .FirstOrDefault(r => this._approvalSnapshots.ContainsKey(r.RequestId));

        if (approvalResponse is null)
        {
            await this.AssignErrorAsync(context, "No pending approval matched the response.").ConfigureAwait(false);
            return;
        }

        if (!approvalResponse.Approved)
        {
            this._approvalSnapshots.TryRemove(approvalResponse.RequestId, out _);
            await this.AssignErrorAsync(context, "MCP tool invocation was not approved by user.").ConfigureAwait(false);
            return;
        }

        // Source invocation fields from the snapshot captured at approval-request time.
        // Headers are re-evaluated (they may contain auth secrets not persisted to state).
        if (!this._approvalSnapshots.TryRemove(approvalResponse.RequestId, out ApprovalSnapshot? snapshot))
        {
            await this.AssignErrorAsync(context, "No pending approval matched the response.").ConfigureAwait(false);
            return;
        }

        Dictionary<string, string>? headers = this.GetHeaders();

        McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync(
            snapshot.ServerUrl,
            snapshot.ServerLabel,
            snapshot.ToolName,
            snapshot.Arguments,
            headers,
            snapshot.ConnectionName,
            cancellationToken).ConfigureAwait(false);

        await this.ProcessResultAsync(context, resultContent, cancellationToken).ConfigureAwait(false);
    }

    /// <summary>
    /// Completes the MCP tool invocation by raising the completion event.
    /// </summary>
    public async ValueTask CompleteAsync(IWorkflowContext context, ActionExecutorResult message, CancellationToken cancellationToken)
    {
        await context.RaiseCompletionEventAsync(this.Model, cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc/>
    public override ValueTask ResetAsync()
    {
        this._approvalSnapshots.Clear();
        return default;
    }

    /// <inheritdoc/>
    /// <remarks>
    /// Persists pending approval snapshots to workflow state so they survive
    /// checkpoint/restore cycles.
    /// </remarks>
    protected override async ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
    {
        Dictionary<string, ApprovalSnapshot> snapshotCopy = this._approvalSnapshots.ToDictionary(kvp => kvp.Key, kvp => kvp.Value, StringComparer.Ordinal);
        await context.QueueStateUpdateAsync(ApprovalSnapshotStateKey, snapshotCopy, null, cancellationToken).ConfigureAwait(false);
        await base.OnCheckpointingAsync(context, cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc/>
    /// <remarks>
    /// Restores pending approval snapshots from workflow state after a checkpoint restore.
    /// </remarks>
    protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
    {
        await base.OnCheckpointRestoredAsync(context, cancellationToken).ConfigureAwait(false);

        this._approvalSnapshots.Clear();
        Dictionary<string, ApprovalSnapshot>? snapshots = await context.ReadStateAsync<Dictionary<string, ApprovalSnapshot>>(
            ApprovalSnapshotStateKey, null, cancellationToken).ConfigureAwait(false);
        if (snapshots is not null)
        {
            foreach (KeyValuePair<string, ApprovalSnapshot> entry in snapshots)
            {
                this._approvalSnapshots[entry.Key] = entry.Value;
            }
        }

        // Migrate a single ApprovalSnapshot at the legacy key under this.Id so the
        // legacy approval response matches the per-invocation map; clear the legacy key.
        ApprovalSnapshot? legacy = await context.ReadStateAsync<ApprovalSnapshot>(
            LegacyApprovalSnapshotStateKey, null, cancellationToken).ConfigureAwait(false);
        if (legacy is not null)
        {
            this._approvalSnapshots.TryAdd(this.Id, legacy);
            await context.QueueStateUpdateAsync<ApprovalSnapshot?>(
                LegacyApprovalSnapshotStateKey, null, null, cancellationToken).ConfigureAwait(false);
        }
    }

    private async ValueTask ProcessResultAsync(IWorkflowContext context, McpServerToolResultContent resultContent, CancellationToken cancellationToken)
    {
        bool autoSend = this.GetAutoSendValue();
        string? conversationId = this.GetConversationId();

        await this.AssignResultAsync(context, resultContent).ConfigureAwait(false);
        ChatMessage resultMessage = new(ChatRole.Tool, resultContent.Outputs);

        // Store messages if output path is configured
        if (this.Model.Output?.Messages is not null)
        {
            await this.AssignAsync(this.Model.Output.Messages?.Path, resultMessage.ToFormula(), context).ConfigureAwait(false);
        }

        // Auto-send the result if configured
        if (autoSend)
        {
            AgentResponse resultResponse = new([resultMessage]);
            await context.AddEventAsync(new AgentResponseEvent(this.Id, resultResponse), cancellationToken).ConfigureAwait(false);
        }

        // Add messages to conversation if conversationId is provided
        if (conversationId is not null)
        {
            ChatMessage assistantMessage = new(ChatRole.Assistant, resultContent.Outputs);
            await agentProvider.CreateMessageAsync(conversationId, assistantMessage, cancellationToken).ConfigureAwait(false);
        }
    }

    private async ValueTask AssignResultAsync(IWorkflowContext context, McpServerToolResultContent toolResult)
    {
        if (this.Model.Output?.Result is null || toolResult.Outputs is null || toolResult.Outputs.Count == 0)
        {
            return;
        }

        List<object?> parsedResults = [];
        foreach (AIContent resultContent in toolResult.Outputs)
        {
            object? resultValue = resultContent switch
            {
                TextContent text => text.Text,
                DataContent data => data.Uri,
                _ => resultContent.ToString(),
            };

            // Convert JsonElement to its raw JSON string for processing
            if (resultValue is JsonElement jsonElement)
            {
                resultValue = jsonElement.GetRawText();
            }

            // Attempt to parse as JSON if it's a string (or was converted from JsonElement)
            if (resultValue is string jsonString)
            {
                try
                {
                    using JsonDocument jsonDocument = JsonDocument.Parse(jsonString);

                    // Handle different JSON value kinds
                    object? parsedValue = jsonDocument.ParseJsonValue(jsonString);

                    parsedResults.Add(parsedValue);
                    continue;
                }
                catch (JsonException)
                {
                    // Not a valid JSON
                }
            }

            parsedResults.Add(resultValue);
        }

        await this.AssignAsync(this.Model.Output.Result?.Path, parsedResults.ToFormula(), context).ConfigureAwait(false);
    }

    private async ValueTask AssignErrorAsync(IWorkflowContext context, string errorMessage)
    {
        // Store error in result if configured (as a simple string)
        if (this.Model.Output?.Result is not null)
        {
            await this.AssignAsync(this.Model.Output.Result?.Path, $"Error: {errorMessage}".ToFormula(), context).ConfigureAwait(false);
        }
    }

    private string GetServerUrl() =>
        this.Evaluator.GetValue(
            Throw.IfNull(
                this.Model.ServerUrl,
                $"{nameof(this.Model)}.{nameof(this.Model.ServerUrl)}")).Value;

    private string? GetServerLabel()
    {
        if (this.Model.ServerLabel is null)
        {
            return null;
        }

        string value = this.Evaluator.GetValue(this.Model.ServerLabel).Value;
        return value.Length == 0 ? null : value;
    }

    private string GetToolName() =>
        this.Evaluator.GetValue(
            Throw.IfNull(
                this.Model.ToolName,
                $"{nameof(this.Model)}.{nameof(this.Model.ToolName)}")).Value;

    private string? GetConversationId()
    {
        if (this.Model.ConversationId is null)
        {
            return null;
        }

        string value = this.Evaluator.GetValue(this.Model.ConversationId).Value;
        return value.Length == 0 ? null : value;
    }

    private bool GetRequireApproval()
    {
        if (this.Model.RequireApproval is null)
        {
            return false;
        }

        return this.Evaluator.GetValue(this.Model.RequireApproval).Value;
    }

    private bool GetAutoSendValue()
    {
        // InvokeToolOutput.AutoSend is never null — it returns a literal-false default
        // when the YAML omits the field. Use AutoSendIsDefaultValue to distinguish an
        // explicit autoSend value from the implicit default, and treat the implicit
        // default as autoSend = true (the historical behavior).
        if (this.Model.Output is { AutoSendIsDefaultValue: false } output)
        {
            return this.Evaluator.GetValue(output.AutoSend).Value;
        }

        return true;
    }

    private string? GetConnectionName()
    {
        if (this.Model.Connection?.Name is null)
        {
            return null;
        }

        string value = this.Evaluator.GetValue(this.Model.Connection.Name).Value;
        return value.Length == 0 ? null : value;
    }

    private Dictionary<string, object?>? GetArguments()
    {
        if (this.Model.Arguments is null)
        {
            return null;
        }

        Dictionary<string, object?> result = [];
        foreach (KeyValuePair<string, ValueExpression> argument in this.Model.Arguments)
        {
            result[argument.Key] = this.Evaluator.GetValue(argument.Value).Value.ToObject();
        }

        return result;
    }

    private Dictionary<string, string>? GetHeaders()
    {
        if (this.Model.Headers is null)
        {
            return null;
        }

        Dictionary<string, string> result = [];
        foreach (KeyValuePair<string, StringExpression> header in this.Model.Headers)
        {
            string value = this.Evaluator.GetValue(header.Value).Value;
            if (!string.IsNullOrEmpty(value))
            {
                result[header.Key] = value;
            }
        }

        return result;
    }

    /// <summary>
    /// Captured invocation parameters used by <see cref="CaptureResponseAsync"/> on
    /// resume so the approved values are invoked regardless of subsequent state changes.
    /// </summary>
    internal sealed record ApprovalSnapshot(
        string ServerUrl,
        string? ServerLabel,
        string ToolName,
        Dictionary<string, object?>? Arguments,
        string? ConnectionName);
}
