// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using Microsoft.Build.Framework;
using Microsoft.Build.Utilities;
using System;
using System.IO;
using System.Reflection.Metadata;
using System.Reflection.Metadata.Ecma335;
using System.Reflection.PortableExecutable;

namespace Microsoft.DotNet.GenFacades;

/// <summary>
/// Rewrites an Assembly's references to be version 0.0.0.0.
/// </summary>
public class ClearAssemblyReferenceVersions : Task
{
    /// <summary>
    /// Assembly to rewrite.
    /// </summary>
    [Required]
    public string Assembly { get; set; }
    
    public override bool Execute()
    {
        try
        {
            using (FileStream stream = File.Open(Assembly, FileMode.Open, FileAccess.ReadWrite, FileShare.Read))
            using (PEReader peReader = new PEReader(stream))
            {
                using (BinaryWriter writer = new BinaryWriter(stream))
                {
                    MetadataReader mdReader = peReader.GetMetadataReader();
                    int assemblyRefOffset = mdReader.GetTableMetadataOffset(TableIndex.AssemblyRef);
                    int numAssemblyRef = mdReader.GetTableRowCount(TableIndex.AssemblyRef);
                    int sizeAssemblyRefRow = mdReader.GetTableRowSize(TableIndex.AssemblyRef);

                    for (int assemblyRefRow = 0; assemblyRefRow < numAssemblyRef; assemblyRefRow++)
                    {
                        stream.Position = peReader.PEHeaders.MetadataStartOffset + assemblyRefOffset + assemblyRefRow * sizeAssemblyRefRow;
                        // see ECMA-335 II.22.5 
                        // row starts with
                        // - MajorVersion, MinorVersion, BuildNumber, RevisionNumber (each being 2-byte constants)
                        // 8 bytes total -> long
                        writer.Write((long)0);
                    }

                    // remove signature, if present, since we changed the binary which would break any signature.
                    // adapted from https://github.com/dotnet/arcade/blob/866b2acd5ffc3c2031c102f6415fd7c6a1a370d5/src/Microsoft.DotNet.Arcade.Sdk/src/Unsign.cs#L58-L65
                    if (peReader.PEHeaders.PEHeader.CertificateTableDirectory.Size != 0)
                    {
                        // see https://docs.microsoft.com/en-us/windows/win32/debug/pe-format#optional-header-data-directories-image-only 
                        int certificateTableDirectoryOffset = (peReader.PEHeaders.PEHeader.Magic == PEMagic.PE32Plus) ? 144 : 128;
                        stream.Position = peReader.PEHeaders.PEHeaderStartOffset + certificateTableDirectoryOffset;

                        writer.Write((long)0);
                    }
                }
            }
        }
        catch (Exception e)
        {
            Log.LogErrorFromException(e, showStackTrace: false);
        }

        return !Log.HasLoggedErrors;
    }
}