// 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.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.ErrorReporting;
using Microsoft.CodeAnalysis.Internal.Log;
using Microsoft.CodeAnalysis.Remote;
using Microsoft.CodeAnalysis.Serialization;
using Roslyn.Utilities;

namespace Microsoft.CodeAnalysis;

internal sealed partial class SolutionCompilationState
{
    /// <summary>
    /// Checksum representing the full checksum tree for this solution compilation state.  Includes the checksum for
    /// <see cref="SolutionState"/>, as well as the checksums for <see cref="FrozenSourceGeneratedDocumentStates"/>
    /// if present.
    /// </summary>
    private readonly AsyncLazy<SolutionCompilationStateChecksums> _lazyChecksums;

    /// <summary>
    /// Mapping from project-id to the checksums needed to synchronize it over to an OOP host.  Lock this specific
    /// field before reading/writing to it.
    /// </summary>
    private readonly Dictionary<ProjectId, AsyncLazy<(SolutionCompilationStateChecksums checksums, ProjectCone projectCone)>> _lazyProjectChecksums = [];

    public bool TryGetStateChecksums([NotNullWhen(true)] out SolutionCompilationStateChecksums? stateChecksums)
        => _lazyChecksums.TryGetValue(out stateChecksums);

    public bool TryGetStateChecksums(ProjectId projectId, [NotNullWhen(true)] out SolutionCompilationStateChecksums? stateChecksums)
    {
        AsyncLazy<(SolutionCompilationStateChecksums checksums, ProjectCone projectCone)>? lazyChecksums;
        lock (_lazyProjectChecksums)
        {
            if (!_lazyProjectChecksums.TryGetValue(projectId, out lazyChecksums) ||
                lazyChecksums == null)
            {
                stateChecksums = null;
                return false;
            }
        }

        if (!lazyChecksums.TryGetValue(out var checksumsAndCone))
        {
            stateChecksums = null;
            return false;
        }

        stateChecksums = checksumsAndCone.checksums;
        return true;
    }

    public Task<SolutionCompilationStateChecksums> GetStateChecksumsAsync(CancellationToken cancellationToken)
        => _lazyChecksums.GetValueAsync(cancellationToken);

    public async Task<Checksum> GetChecksumAsync(CancellationToken cancellationToken)
    {
        var collection = await GetStateChecksumsAsync(cancellationToken).ConfigureAwait(false);
        return collection.Checksum;
    }

    /// <summary>Gets the checksum for only the requested project (and any project it depends on)</summary>
    public async Task<(SolutionCompilationStateChecksums checksums, ProjectCone projectCone)> GetStateChecksumsAsync(
        ProjectId projectId,
        CancellationToken cancellationToken)
    {
        Contract.ThrowIfNull(projectId);

        AsyncLazy<(SolutionCompilationStateChecksums checksums, ProjectCone projectCone)>? checksums;
        lock (_lazyProjectChecksums)
        {
            if (!_lazyProjectChecksums.TryGetValue(projectId, out checksums))
            {
                checksums = AsyncLazy.Create(static async (arg, cancellationToken) =>
                {
                    var (checksum, projectCone) = await arg.self.ComputeChecksumsAsync(arg.projectId, cancellationToken).ConfigureAwait(false);
                    Contract.ThrowIfNull(projectCone);
                    return (checksum, projectCone);
                }, arg: (self: this, projectId));

                _lazyProjectChecksums.Add(projectId, checksums);
            }
        }

        var collection = await checksums.GetValueAsync(cancellationToken).ConfigureAwait(false);
        return collection;
    }

    /// <summary>Gets the checksum for only the requested project (and any project it depends on)</summary>
    public async Task<Checksum> GetChecksumAsync(ProjectId projectId, CancellationToken cancellationToken)
    {
        var (checksums, _) = await GetStateChecksumsAsync(projectId, cancellationToken).ConfigureAwait(false);
        return checksums.Checksum;
    }

    private async Task<(SolutionCompilationStateChecksums checksums, ProjectCone? projectCone)> ComputeChecksumsAsync(
        ProjectId? projectId,
        CancellationToken cancellationToken)
    {
        try
        {
            using (Logger.LogBlock(FunctionId.SolutionCompilationState_ComputeChecksumsAsync, this.SolutionState.FilePath, cancellationToken))
            {
                Checksum solutionStateChecksum;
                ProjectCone? projectCone;

                if (projectId is null)
                {
                    solutionStateChecksum = await this.SolutionState.GetChecksumAsync(cancellationToken).ConfigureAwait(false);
                    projectCone = null;
                }
                else
                {
                    var stateChecksums = await this.SolutionState.GetStateChecksumsAsync(projectId, cancellationToken).ConfigureAwait(false);
                    solutionStateChecksum = stateChecksums.Checksum;
                    projectCone = stateChecksums.ProjectCone;
                }

                var serializer = this.SolutionState.Services.GetRequiredService<ISerializerService>();
                var identityChecksums = FrozenSourceGeneratedDocumentStates.SelectAsArray(
                    static (s, arg) => arg.serializer.CreateChecksum(s.Identity, cancellationToken: arg.cancellationToken), (serializer, cancellationToken));

                var frozenSourceGeneratedDocumentTexts = await FrozenSourceGeneratedDocumentStates.GetDocumentChecksumsAndIdsAsync(cancellationToken).ConfigureAwait(false);
                var frozenSourceGeneratedDocumentIdentities = new ChecksumCollection(identityChecksums);
                var frozenSourceGeneratedDocumentGenerationDateTimes = FrozenSourceGeneratedDocumentStates.SelectAsArray(d => d.GenerationDateTime);

                // Ensure we only send the execution map over for projects in the project cone.
                var versionMapChecksum = this.GetFilteredSourceGenerationExecutionMap(projectCone).GetChecksum();

                var compilationStateChecksums = new SolutionCompilationStateChecksums(
                    solutionStateChecksum,
                    versionMapChecksum,
                    frozenSourceGeneratedDocumentTexts,
                    frozenSourceGeneratedDocumentIdentities,
                    frozenSourceGeneratedDocumentGenerationDateTimes);
                return (compilationStateChecksums, projectCone);
            }
        }
        catch (Exception e) when (FatalError.ReportAndPropagateUnlessCanceled(e, cancellationToken))
        {
            throw ExceptionUtilities.Unreachable();
        }
    }

    public SourceGeneratorExecutionVersionMap GetFilteredSourceGenerationExecutionMap(ProjectCone? projectCone)
    {
        var builder = this.SourceGeneratorExecutionVersionMap.Map.ToBuilder();

        foreach (var projectState in this.SolutionState.SortedProjectStates)
        {
            var projectId = projectState.Id;
            if (!RemoteSupportedLanguages.IsSupported(projectState.Language))
            {
                builder.Remove(projectId);
            }
            else if (projectCone != null && !projectCone.Contains(projectId))
            {
                builder.Remove(projectId);
            }
        }

        if (builder.Count == this.SourceGeneratorExecutionVersionMap.Map.Count)
            return this.SourceGeneratorExecutionVersionMap;

        return new SourceGeneratorExecutionVersionMap(builder.ToImmutable());
    }
}