var builder = DistributedApplication.CreateBuilder(args);

var cache = builder.AddRedis("cache")
    .WithArchitectureDiagramNodeDetails("Redis cache");

var apiService = builder.AddProject<Projects.AspireDiagrams02_ApiService>("apiservice")
    .WithArchitectureDiagramNodeDetails("API service");

builder.AddProject<Projects.AspireDiagrams02_Web>("webfrontend")
    .WithArchitectureDiagramNodeDetails("Web frontend")
    .WithExternalHttpEndpoints()
    .WithReference(cache).WaitFor(cache)
    .WithReference(apiService).WaitFor(apiService);

builder.GenerateArchitectureDiagramOnPublish(options =>
{
    options.Title = "Aspire starter with Redis";
});

builder.Build().Run();

using System.Diagnostics;
using System.Text;
using Aspire.Hosting.Lifecycle;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Options;

namespace Aspire.Hosting;

// This is a somewhat basic implementation to demonstrate the concept, and it can be extended in many ways.
// CONSIDER support for grouping nodes into subgraphs
// CONSIDER support for arbitrary text (e.g. a floating comment)
// CONSIDER support for diagram configuration (layout, look, theme)
// CONSIDER support for custom node shape, color, and style
// CONSIDER support for custom relationship label, arrow style and line style
// CONSIDER modeling the diagram as a graph to support different rendering strategies
internal static class ArchitectureDiagramExtensions
{
    public static IDistributedApplicationBuilder GenerateArchitectureDiagramOnPublish(this IDistributedApplicationBuilder builder, Action<ArchitectureDiagramOptions>? configure = null)
    {
        if (configure != null)
        {
            builder.Services.Configure(configure);
        }

        builder.Services.TryAddLifecycleHook<ArchitectureDiagramGeneratorLifecycleHook>();

        return builder;
    }

    public static IResourceBuilder<T> WithArchitectureDiagramNodeDetails<T>(this IResourceBuilder<T> builder, string label)
        where T : IResource
    {
        ArgumentException.ThrowIfNullOrEmpty(label);
        return builder.WithAnnotation(new ArchitectureDiagramNodeDetailsAnnotation(label), ResourceAnnotationMutationBehavior.Replace);
    }

    public static IResourceBuilder<T> ExcludeFromArchitectureDiagram<T>(this IResourceBuilder<T> builder)
        where T : IResource
    {
        return builder.WithAnnotation(new ExcludeFromArchitectureDiagramAnnotation(), ResourceAnnotationMutationBehavior.Replace);
    }
}

internal sealed class ArchitectureDiagramGeneratorLifecycleHook(
    DistributedApplicationExecutionContext executionContext,
    IOptions<ArchitectureDiagramOptions> diagramOptions) : IDistributedApplicationLifecycleHook
{
    private readonly StringBuilder _mermaidOutput = new();
    private readonly HashSet<IResource> _visibleResources = [];
    private readonly HashSet<(IResource Source, IResource Destination)> _visitedRelationships = [];

    public async Task BeforeStartAsync(DistributedApplicationModel appModel, CancellationToken cancellationToken = default)
    {
        if (!executionContext.IsPublishMode || appModel.Resources.Count == 0)
        {
            return;
        }

        // Ordering by name increases chances of git diffs being more readable
        var orderedResources = appModel.Resources.OrderBy(x => x.Name, StringComparer.Ordinal).ToArray();

        this.AppendDiagramTitle();
        this.AppendDiagramDirection();
        this.AppendResourceNodes(orderedResources);
        this.AppendResourcesRelationships(orderedResources);

        // CONSIDER specifying the file name/path from the options or CLI args with configuration binding (same as manifest's --output-path)
        // CONSIDER resolving the containing directory path from CLI args, content root path or project directory
        await File.WriteAllTextAsync("architecture.mmd", this._mermaidOutput.ToString(), cancellationToken).ConfigureAwait(false);
    }

    private void AppendDiagramTitle()
    {
        if (!string.IsNullOrEmpty(diagramOptions.Value.Title))
        {
            this._mermaidOutput.AppendLine("---");
            this._mermaidOutput.Append("title: ").AppendLine(diagramOptions.Value.Title);
            this._mermaidOutput.AppendLine("---");
        }
    }

    private void AppendDiagramDirection()
    {
        var direction = diagramOptions.Value.Direction switch
        {
            DiagramDirection.TopToBottom => "TB",
            DiagramDirection.BottomToTop => "BT",
            DiagramDirection.LeftToRight => "LR",
            DiagramDirection.RightToLeft => "RL",
            _ => throw new ArgumentOutOfRangeException(nameof(diagramOptions.Value.Direction), diagramOptions.Value.Direction, "Invalid diagram direction")
        };

        this._mermaidOutput.Append("flowchart ").AppendLine(direction);
    }

    private void AppendResourceNodes(IResource[] resources)
    {
        foreach (var resource in resources)
        {
            AppendResourceNode(resource);
        }

        this._mermaidOutput.AppendLine();
    }

    private void AppendResourceNode(IResource resource)
    {
        var isExcluded = resource.Annotations.OfType<ExcludeFromArchitectureDiagramAnnotation>().Any() ||
            resource.Annotations.OfType<ResourceSnapshotAnnotation>().FirstOrDefault()?.InitialSnapshot.State == KnownResourceStates.Hidden;

        if (isExcluded)
        {
            return;
        }

        var nodeDetails = resource.Annotations.OfType<ArchitectureDiagramNodeDetailsAnnotation>().FirstOrDefault()
            ?? new ArchitectureDiagramNodeDetailsAnnotation(resource.Name);

        this._visibleResources.Add(resource);
        this._mermaidOutput.Append("    ").Append(resource.Name).Append('(').Append(nodeDetails.Label).AppendLine(")");
    }

    private void AppendResourcesRelationships(IResource[] resources)
    {
        foreach (var resource in resources)
        {
            var relationships = resource.Annotations.OfType<ResourceRelationshipAnnotation>();

            foreach (var relationship in relationships)
            {
                this.AppendResourceRelationships(resource, relationship);
            }
        }

        this._mermaidOutput.AppendLine();
    }

    // https://github.com/dotnet/aspire/blob/39b8606f1b43ae4bd4ef5ea3492f0fef7fe548b2/src/Shared/Model/KnownRelationshipTypes.cs#L8
    private const string WaitForKnownRelationshipType = "WaitFor";

    private void AppendResourceRelationships(IResource resource, ResourceRelationshipAnnotation relationship)
    {
        if (relationship.Type == WaitForKnownRelationshipType)
        {
            return;
        }

        var isSourceExcluded = !this._visibleResources.Contains(resource);
        var isTargetExcluded = !this._visibleResources.Contains(relationship.Resource);
        var isAlreadyVisited = this._visitedRelationships.Contains((resource, relationship.Resource));

        if (isSourceExcluded || isTargetExcluded || isAlreadyVisited)
        {
            return;
        }

        // CONSIDER grouping relationships from the same source to the same target into a single one
        this._visitedRelationships.Add((resource, relationship.Resource));

        // CONSIDER showing the relationship type if different than the built-in "Reference" type
        this._mermaidOutput.Append("    ").Append(resource.Name).Append(" --> ").AppendLine(relationship.Resource.Name);
    }
}

internal sealed class ArchitectureDiagramOptions
{
    public string? Title { get; set; }

    public DiagramDirection Direction { get; set; } = DiagramDirection.TopToBottom;
}

[DebuggerDisplay("Type = {GetType().Name,nq}, Label = {Label}")]
internal sealed class ArchitectureDiagramNodeDetailsAnnotation(string label) : IResourceAnnotation
{
    public string Label { get; } = label;
}

[DebuggerDisplay("Type = {GetType().Name,nq}")]
internal sealed class ExcludeFromArchitectureDiagramAnnotation : IResourceAnnotation;

internal enum DiagramDirection
{
    TopToBottom,
    BottomToTop,
    LeftToRight,
    RightToLeft,
}