// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.CommandLine;
using System.Globalization;
using System.Text.Json;
using Aspire.Cli.Backchannel;
using Aspire.Cli.Resources;
using ModelContextProtocol.Protocol;
 
namespace Aspire.Cli.Commands;
 
/// <summary>
/// Calls an MCP tool on a running Aspire resource.
/// </summary>
internal sealed class McpCallCommand : BaseCommand
{
    internal override HelpGroup HelpGroup => HelpGroup.ToolsAndConfiguration;
 
    private readonly AppHostConnectionResolver _connectionResolver;
 
    private static readonly Argument<string> s_resourceArgument = new("resource")
    {
        Description = McpCommandStrings.CallCommand_ResourceArgumentDescription
    };
 
    private static readonly Argument<string> s_toolArgument = new("tool")
    {
        Description = McpCommandStrings.CallCommand_ToolArgumentDescription
    };
 
    private static readonly Option<string?> s_inputOption = new("--input", "-i")
    {
        Description = McpCommandStrings.CallCommand_InputOptionDescription
    };
 
    private static readonly OptionWithLegacy<FileInfo?> s_appHostOption = new("--apphost", "--project", SharedCommandStrings.AppHostOptionDescription);
 
    public McpCallCommand(
        AppHostConnectionResolver connectionResolver,
        CommonCommandServices services)
        : base("call", McpCommandStrings.CallCommand_Description, services)
    {
        _connectionResolver = connectionResolver;
 
        Arguments.Add(s_resourceArgument);
        Arguments.Add(s_toolArgument);
        Options.Add(s_inputOption);
        Options.Add(s_appHostOption);
    }
 
    protected override async Task<CommandResult> ExecuteAsync(ParseResult parseResult, CancellationToken cancellationToken)
    {
        var resourceName = parseResult.GetValue(s_resourceArgument)!;
        var toolName = parseResult.GetValue(s_toolArgument)!;
        var inputJson = parseResult.GetValue(s_inputOption);
        var passedAppHostProjectFile = parseResult.GetValue(s_appHostOption);
 
        var result = await _connectionResolver.ResolveConnectionAsync(
            passedAppHostProjectFile,
            SharedCommandStrings.ScanningForRunningAppHosts,
            string.Format(CultureInfo.CurrentCulture, SharedCommandStrings.SelectAppHost, "call MCP tool on"),
            SharedCommandStrings.AppHostNotRunning,
            cancellationToken);
 
        if (!result.Success)
        {
            return CommandResult.FromExitCode(AppHostConnectionResultHandler.DisplayFailureAsError(result, InteractionService, CliExitCodes.FailedToDotnetRunAppHost));
        }
 
        var connection = result.Connection!;
 
        // Parse input JSON into arguments dictionary
        IReadOnlyDictionary<string, JsonElement>? arguments = null;
        if (!string.IsNullOrEmpty(inputJson))
        {
            try
            {
                using var doc = JsonDocument.Parse(inputJson);
                if (doc.RootElement.ValueKind != JsonValueKind.Object)
                {
                    return CommandResult.Failure(CliExitCodes.InvalidCommand, "Invalid JSON input: expected a JSON object.");
                }
                var dict = new Dictionary<string, JsonElement>();
                foreach (var prop in doc.RootElement.EnumerateObject())
                {
                    dict[prop.Name] = prop.Value.Clone();
                }
                arguments = dict;
            }
            catch (JsonException ex)
            {
                return CommandResult.Failure(CliExitCodes.InvalidCommand, $"Invalid JSON input: {ex.Message}");
            }
        }
 
        try
        {
            var callResult = await connection.CallResourceMcpToolAsync(
                resourceName,
                toolName,
                arguments,
                cancellationToken);
 
            // Output the result content
            if (callResult.Content is { Count: > 0 })
            {
                foreach (var content in callResult.Content)
                {
                    if (content is TextContentBlock textContent)
                    {
                        InteractionService.DisplayRawText(textContent.Text);
                    }
                    else
                    {
                        using var stream = new MemoryStream();
                        using (var writer = new Utf8JsonWriter(stream, new JsonWriterOptions { Indented = true }))
                        {
                            writer.WriteStartObject();
                            writer.WriteString("type", content.Type);
                            writer.WriteEndObject();
                        }
                        InteractionService.DisplayRawText(System.Text.Encoding.UTF8.GetString(stream.ToArray()));
                    }
                }
            }
 
            if (callResult.IsError == true)
            {
                return CommandResult.Failure(CliExitCodes.InvalidCommand);
            }
 
            return CommandResult.Success();
        }
        catch (Exception ex)
        {
            return CommandResult.Failure(CliExitCodes.InvalidCommand, $"Failed to call tool '{toolName}' on resource '{resourceName}': {ex.Message}");
        }
    }
}