// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using Microsoft.Testing.Platform.CommandLine;
using Microsoft.Testing.Platform.Extensions.OutputDevice;
using Microsoft.Testing.Platform.Helpers;
using Microsoft.Testing.Platform.IPC.Models;
using Microsoft.Testing.Platform.OutputDevice;
using Microsoft.Testing.Platform.Resources;
using Microsoft.Testing.Platform.ServerMode;
using Microsoft.Testing.Platform.Tools;

namespace Microsoft.Testing.Platform.Extensions.ArtifactPostProcessing;

internal sealed class ArtifactPostProcessingDispatcherTool(
    ICommandLineOptions commandLineOptions,
    IReadOnlyList<IArtifactPostProcessor> processors,
    IPushOnlyProtocol? protocol,
    IEnvironment environment,
    IOutputDevice outputDevice) : ITool, IOutputDeviceDataProducer
{
    public const string ToolName = "internal-merge-artifacts";

    private readonly ICommandLineOptions _commandLineOptions = commandLineOptions;
    private readonly IEnvironment _environment = environment;
    private readonly IOutputDevice _outputDevice = outputDevice;
    private readonly IReadOnlyList<IArtifactPostProcessor> _processors = processors;
    private readonly IPushOnlyProtocol? _protocol = protocol;

    public string Name => ToolName;

    public string Uid => "Microsoft.Testing.Platform.ArtifactPostProcessing.Dispatcher";

    public string Version => PlatformVersion.Version;

    public string DisplayName => PlatformResources.ArtifactPostProcessingDispatcherDisplayName;

    public string Description => PlatformResources.ArtifactPostProcessingDispatcherDescription;

    public Task<bool> IsEnabledAsync() => Task.FromResult(true);

    public async Task<int> RunAsync(CancellationToken cancellationToken)
    {
        if (_protocol is null || !_protocol.IsServerMode)
        {
            await DisplayErrorAsync(PlatformResources.ArtifactPostProcessingDispatcherRequiresPipe, cancellationToken).ConfigureAwait(false);
            return (int)ExitCode.InvalidCommandLine;
        }

        ApplicationStateGuard.Ensure(
            _commandLineOptions.TryGetOptionArgumentList(ArtifactPostProcessingDispatcherToolCommandLine.ManifestOptionName, out string[]? manifestPaths));
        RoslynDebug.Assert(manifestPaths.Length == 1);

        ArtifactPostProcessingManifest manifest;
        try
        {
            manifest = ArtifactPostProcessingManifest.Load(manifestPaths[0]);
            Directory.CreateDirectory(manifest.OutputDirectory);
        }
        catch (Exception ex) when (ex is IOException or UnauthorizedAccessException or FormatException or ArgumentException or NotSupportedException)
        {
            await DisplayErrorAsync(
                string.Format(CultureInfo.CurrentCulture, PlatformResources.ArtifactPostProcessingManifestReadError, ex.Message),
                cancellationToken).ConfigureAwait(false);
            return (int)ExitCode.InvalidCommandLine;
        }

        IArtifactPostProcessor[] eligibleProcessors =
        [
            .. _processors.Where(processor =>
                processor.SupportedModes.Contains(manifest.Context.Mode)
                && (!manifest.Context.IsTruncated || processor.SupportsTruncatedRuns)),
        ];
        await WarnAboutProcessorConflictsAsync(eligibleProcessors, cancellationToken).ConfigureAwait(false);

        List<InputArtifact> unmatchedInputs = [.. manifest.Inputs];
        List<FileArtifactMessage> outputs = [];
        bool failed = false;
        foreach (IArtifactPostProcessor processor in eligibleProcessors)
        {
            InputArtifact[] matchingInputs = [.. unmatchedInputs.Where(input => Matches(processor, input))];
            if (matchingInputs.Length == 0)
            {
                continue;
            }

            unmatchedInputs.RemoveAll(input => matchingInputs.Contains(input));
            try
            {
                if (await processor.ProcessAsync(
                    matchingInputs,
                    manifest.OutputDirectory,
                    manifest.Context,
                    cancellationToken).ConfigureAwait(false) is { } output)
                {
                    ProcessedArtifact validatedOutput = ValidateProcessedArtifact(output, manifest.OutputDirectory, matchingInputs);
                    outputs.Add(new FileArtifactMessage(
                        validatedOutput.Path,
                        validatedOutput.DisplayName,
                        validatedOutput.Description,
                        TestUid: null,
                        TestDisplayName: null,
                        SessionUid: null,
                        validatedOutput.Kind,
                        [.. matchingInputs.Select(input => input.Path)]));
                }
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
            {
                throw;
            }
            catch (Exception ex)
            {
                failed = true;
                await DisplayWarningAsync(
                    string.Format(CultureInfo.CurrentCulture, PlatformResources.ArtifactPostProcessingDispatcherProcessorFailed, processor.Uid, ex.Message),
                    cancellationToken).ConfigureAwait(false);
            }
        }

        if (outputs.Count > 0)
        {
            await _protocol.SendMessageAsync(new FileArtifactMessages(
                _environment.GetEnvironmentVariable(EnvironmentVariableConstants.TESTINGPLATFORM_DOTNETTEST_EXECUTIONID),
                _protocol.InstanceId,
                [.. outputs])).ConfigureAwait(false);
        }

        return failed ? (int)ExitCode.GenericFailure : (int)ExitCode.Success;
    }

    internal static ProcessedArtifact ValidateProcessedArtifact(
        ProcessedArtifact output,
        string outputDirectory,
        IReadOnlyList<InputArtifact> inputs)
    {
        string outputPath = Path.GetFullPath(output.Path);
        string outputDirectoryPrefix = Path.GetFullPath(outputDirectory)
            .TrimEnd(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar)
            + Path.DirectorySeparatorChar;
        StringComparison pathComparison = RuntimeInformation.IsOSPlatform(OSPlatform.Windows)
            ? StringComparison.OrdinalIgnoreCase
            : StringComparison.Ordinal;

        return outputPath.StartsWith(outputDirectoryPrefix, pathComparison)
            && File.Exists(outputPath)
            && !inputs.Any(input => string.Equals(Path.GetFullPath(input.Path), outputPath, pathComparison))
                ? new ProcessedArtifact(outputPath, output.Kind, output.DisplayName, output.Description)
                : throw new InvalidOperationException(PlatformResources.ArtifactPostProcessingDispatcherInvalidProcessorOutput);
    }

    private static bool Matches(IArtifactPostProcessor processor, InputArtifact input)
        => input.Kind is not null
            ? processor.SupportedKinds.Contains(input.Kind, StringComparer.Ordinal)
            : processor.SupportedFileExtensionsFallback.Contains(Path.GetExtension(input.Path), StringComparer.OrdinalIgnoreCase);

    private Task DisplayErrorAsync(string message, CancellationToken cancellationToken)
        => _outputDevice.DisplayAsync(this, new ErrorMessageOutputDeviceData(message), cancellationToken);

    private Task DisplayWarningAsync(string message, CancellationToken cancellationToken)
        => _outputDevice.DisplayAsync(this, new WarningMessageOutputDeviceData(message), cancellationToken);

    private async Task WarnAboutProcessorConflictsAsync(
        IReadOnlyList<IArtifactPostProcessor> processors,
        CancellationToken cancellationToken)
    {
        foreach (KeyValuePair<string, string> conflict in FindProcessorConflicts(processors))
        {
            await _outputDevice.DisplayAsync(
                this,
                new WarningMessageOutputDeviceData(string.Format(
                    CultureInfo.CurrentCulture,
                    PlatformResources.ArtifactPostProcessingDispatcherProcessorConflict,
                    conflict.Key,
                    conflict.Value)),
                cancellationToken).ConfigureAwait(false);
        }
    }

    internal static IReadOnlyDictionary<string, string> FindProcessorConflicts(IEnumerable<IArtifactPostProcessor> processors)
    {
        Dictionary<string, string> firstProcessorByKind = [with(StringComparer.Ordinal)];
        Dictionary<string, string> firstProcessorByExtension = [with(StringComparer.Ordinal)];
        Dictionary<string, string> conflicts = [with(StringComparer.Ordinal)];
        foreach (IArtifactPostProcessor processor in processors)
        {
            AddConflicts(new HashSet<string>(processor.SupportedKinds, StringComparer.Ordinal), firstProcessorByKind, processor.Uid);
            AddConflicts(
                new HashSet<string>(
                    processor.SupportedFileExtensionsFallback.Select(extension => extension.ToLowerInvariant()),
                    StringComparer.Ordinal),
                firstProcessorByExtension,
                processor.Uid);
        }

        return conflicts;

        void AddConflicts(
            IEnumerable<string> capabilities,
            Dictionary<string, string> firstProcessorByCapability,
            string processorUid)
        {
            foreach (string capability in capabilities)
            {
                if (!firstProcessorByCapability.TryGetValue(capability, out string? firstProcessorUid))
                {
                    firstProcessorByCapability.Add(capability, processorUid);
                }
                else if (firstProcessorUid != processorUid && !conflicts.ContainsKey(capability))
                {
                    conflicts.Add(capability, firstProcessorUid);
                }
            }
        }
    }
}