// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
 
using System;
using System.Collections.Immutable;
using System.Composition;
using System.IO;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.Host.Mef;
using Microsoft.CodeAnalysis.PooledObjects;
using Microsoft.CodeAnalysis.Shared.Extensions;
using Roslyn.Utilities;
 
namespace Microsoft.CodeAnalysis.Host
{
    /// <summary>
    /// Provides a compile-time view of the current workspace solution.
    /// Workaround for Razor projects which generate both design-time and compile-time source files.
    /// TODO: remove https://github.com/dotnet/roslyn/issues/51678
    /// </summary>
    internal sealed class CompileTimeSolutionProvider : ICompileTimeSolutionProvider
    {
        [ExportWorkspaceServiceFactory(typeof(ICompileTimeSolutionProvider), [WorkspaceKind.Host]), Shared]
        private sealed class Factory : IWorkspaceServiceFactory
        {
            [ImportingConstructor]
            [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
            public Factory()
            {
            }
 
            [Obsolete(MefConstruction.FactoryMethodMessage, error: true)]
            public IWorkspaceService? CreateService(HostWorkspaceServices workspaceServices)
                => new CompileTimeSolutionProvider(workspaceServices.Workspace);
        }
 
        private const string RazorEncConfigFileName = "RazorSourceGenerator.razorencconfig";
        private const string RazorSourceGeneratorTypeName = "Microsoft.NET.Sdk.Razor.SourceGenerators.RazorSourceGenerator";
        private static readonly ImmutableArray<string> s_razorSourceGeneratorAssemblyNames =
        [
            "Microsoft.NET.Sdk.Razor.SourceGenerators",
            "Microsoft.CodeAnalysis.Razor.Compiler.SourceGenerators",
            "Microsoft.CodeAnalysis.Razor.Compiler",
        ];
        private static readonly ImmutableArray<string> s_razorSourceGeneratorFileNamePrefixes = s_razorSourceGeneratorAssemblyNames
            .SelectAsArray(static assemblyName => Path.Combine(assemblyName, RazorSourceGeneratorTypeName));
 
        private readonly object _gate = new();
 
        /// <summary>
        /// Cached compile-time solution corresponding to an existing design-time solution.
        /// </summary>
#if NET
        private readonly ConditionalWeakTable<Solution, Solution> _designTimeToCompileTimeSolution = [];
#else
        private ConditionalWeakTable<Solution, Solution> _designTimeToCompileTimeSolution = new();
#endif
 
        private Solution? _lastCompileTimeSolution;
 
        public CompileTimeSolutionProvider(Workspace workspace)
        {
            workspace.WorkspaceChanged += (s, e) =>
            {
                if (e.Kind is WorkspaceChangeKind.SolutionCleared or WorkspaceChangeKind.SolutionRemoved)
                {
                    lock (_gate)
                    {
#if NET
                        _designTimeToCompileTimeSolution.Clear();
#else
                        _designTimeToCompileTimeSolution = new();
#endif
                        _lastCompileTimeSolution = null;
                    }
                }
            };
        }
 
        private static bool IsRazorAnalyzerConfig(TextDocumentState documentState)
            => documentState.FilePath != null && documentState.FilePath.EndsWith(RazorEncConfigFileName, StringComparison.OrdinalIgnoreCase);
 
        public Solution GetCompileTimeSolution(Solution designTimeSolution)
        {
            lock (_gate)
            {
                _designTimeToCompileTimeSolution.TryGetValue(designTimeSolution, out var cachedCompileTimeSolution);
 
                // Design time solution hasn't changed since we calculated the last compile-time solution:
                if (cachedCompileTimeSolution != null)
                    return cachedCompileTimeSolution;
 
                var staleSolution = _lastCompileTimeSolution;
                var compileTimeSolution = designTimeSolution;
 
                foreach (var (_, projectState) in compileTimeSolution.SolutionState.ProjectStates)
                {
                    using var _1 = ArrayBuilder<DocumentId>.GetInstance(out var configIdsToRemove);
                    using var _2 = ArrayBuilder<DocumentId>.GetInstance(out var documentIdsToRemove);
 
                    foreach (var (_, configState) in projectState.AnalyzerConfigDocumentStates.States)
                    {
                        if (IsRazorAnalyzerConfig(configState))
                        {
                            configIdsToRemove.Add(configState.Id);
                        }
                    }
 
                    // only remove design-time only documents when source-generated ones replace them
                    if (configIdsToRemove.Count > 0)
                    {
                        foreach (var (_, documentState) in projectState.DocumentStates.States)
                        {
                            if (documentState.Attributes.DesignTimeOnly || IsRazorDesignTimeDocument(documentState))
                            {
                                documentIdsToRemove.Add(documentState.Id);
                            }
                        }
 
                        compileTimeSolution = compileTimeSolution
                            .RemoveAnalyzerConfigDocuments(configIdsToRemove.ToImmutable())
                            .RemoveDocuments(documentIdsToRemove.ToImmutable());
 
                        if (staleSolution is not null)
                        {
                            var existingStaleProject = staleSolution.GetProject(projectState.Id);
                            if (existingStaleProject is not null)
                                compileTimeSolution = compileTimeSolution.WithCachedSourceGeneratorState(projectState.Id, existingStaleProject);
                        }
                    }
                }
 
                compileTimeSolution = _designTimeToCompileTimeSolution.GetValue(designTimeSolution, _ => compileTimeSolution);
                _lastCompileTimeSolution = compileTimeSolution;
 
                return compileTimeSolution;
            }
        }
 
        // Copied from
        // https://github.com/dotnet/sdk/blob/main/src/RazorSdk/SourceGenerators/RazorSourceGenerator.Helpers.cs#L32
        private static string GetIdentifierFromPath(string filePath)
        {
            var builder = new StringBuilder(filePath.Length);
 
            for (var i = 0; i < filePath.Length; i++)
            {
                switch (filePath[i])
                {
                    case ':' or '\\' or '/':
                    case char ch when !char.IsLetterOrDigit(ch):
                        builder.Append('_');
                        break;
                    default:
                        builder.Append(filePath[i]);
                        break;
                }
            }
 
            return builder.ToString();
        }
 
        private static bool IsRazorDesignTimeDocument(DocumentState documentState)
            => documentState.FilePath?.EndsWith(".razor.g.cs") == true || documentState.FilePath?.EndsWith(".cshtml.g.cs") == true;
 
        internal static async Task<Document?> TryGetCompileTimeDocumentAsync(
            Document designTimeDocument,
            Solution compileTimeSolution,
            CancellationToken cancellationToken,
            string? generatedDocumentPathPrefix = null)
        {
            var compileTimeDocument = await compileTimeSolution.GetDocumentAsync(designTimeDocument.Id, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
            if (compileTimeDocument != null)
            {
                return compileTimeDocument;
            }
 
            if (!IsRazorDesignTimeDocument(designTimeDocument.DocumentState))
            {
                return null;
            }
 
            var designTimeProjectDirectoryName = PathUtilities.GetDirectoryName(designTimeDocument.Project.FilePath)!;
 
            var generatedDocumentPaths = BuildGeneratedDocumentPaths(designTimeProjectDirectoryName, designTimeDocument.FilePath!, generatedDocumentPathPrefix);
 
            var sourceGeneratedDocuments = await compileTimeSolution.GetRequiredProject(designTimeDocument.Project.Id).GetSourceGeneratedDocumentsAsync(cancellationToken).ConfigureAwait(false);
            return sourceGeneratedDocuments.SingleOrDefault(d => d.FilePath != null && generatedDocumentPaths.Contains(d.FilePath));
        }
 
        /// <summary>
        /// Note that in .NET 6 Preview 7 the source generator changed to passing in the relative doc path without a leading \ to GetIdentifierFromPath
        /// which caused the source generated file name to no longer be prefixed by an _.  Additionally, the file extension was changed to .g.cs
        /// </summary>
        private static OneOrMany<string> BuildGeneratedDocumentPaths(string designTimeProjectDirectoryName, string designTimeDocumentFilePath, string? generatedDocumentPathPrefix)
        {
            var relativeDocumentPath = GetRelativeDocumentPath(designTimeProjectDirectoryName, designTimeDocumentFilePath);
 
            if (generatedDocumentPathPrefix is not null)
            {
                return OneOrMany.Create(GetGeneratedDocumentPath(generatedDocumentPathPrefix, relativeDocumentPath));
            }
 
            return OneOrMany.Create(s_razorSourceGeneratorFileNamePrefixes.SelectAsArray(
                static (prefix, relativeDocumentPath) => GetGeneratedDocumentPath(prefix, relativeDocumentPath), relativeDocumentPath));
 
            static string GetGeneratedDocumentPath(string prefix, string relativeDocumentPath)
            {
                return Path.Combine(prefix, GetIdentifierFromPath(relativeDocumentPath)) + ".g.cs";
            }
        }
 
        private static string GetRelativeDocumentPath(string projectDirectory, string designTimeDocumentFilePath)
            => PathUtilities.GetRelativePath(projectDirectory, designTimeDocumentFilePath)[..^".g.cs".Length];
    }
}