// Copyright (c) .NET Foundation and contributors. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
using System.IO;
using System.Text;
 
using Mono.Cecil;
using Mono.Cecil.Cil;
 
namespace Mono.Linker
{
    internal sealed class CustomSymbolWriterProvider : ISymbolWriterProvider
    {
        readonly DefaultSymbolWriterProvider _defaultProvider = new DefaultSymbolWriterProvider();
        readonly bool _preserveSymbolPaths;
 
        public CustomSymbolWriterProvider(bool preserveSymbolPaths) => this._preserveSymbolPaths = preserveSymbolPaths;
 
        public ISymbolWriter GetSymbolWriter(ModuleDefinition module, string fileName)
            => new CustomSymbolWriter(_defaultProvider.GetSymbolWriter(module, fileName), module, _preserveSymbolPaths);
 
        public ISymbolWriter GetSymbolWriter(ModuleDefinition module, Stream symbolStream)
            => new CustomSymbolWriter(_defaultProvider.GetSymbolWriter(module, symbolStream), module, _preserveSymbolPaths);
    }
 
    internal sealed class CustomSymbolWriter : ISymbolWriter
    {
        // ASCII "RSDS": https://github.com/dotnet/runtime/blob/main/docs/design/specs/PE-COFF.md#codeview-debug-directory-entry-type-2
        const int CodeViewSignature = 0x53445352;
 
        readonly ISymbolWriter _symbolWriter;
        readonly ModuleDefinition _module;
        readonly bool _preserveSymbolPaths;
 
        internal CustomSymbolWriter(ISymbolWriter defaultWriter, ModuleDefinition module, bool preserveSymbolPaths)
        {
            _symbolWriter = defaultWriter;
            _module = module;
            _preserveSymbolPaths = preserveSymbolPaths;
        }
 
        public ImageDebugHeader GetDebugHeader()
        {
            var header = _symbolWriter.GetDebugHeader();
            if (!_preserveSymbolPaths)
                return header;
 
            if (!header.HasEntries)
                return header;
 
            for (int i = 0; i < header.Entries.Length; i++)
            {
                header.Entries[i] = ProcessEntry(header.Entries[i]);
            }
 
            return header;
        }
 
        ImageDebugHeaderEntry ProcessEntry(ImageDebugHeaderEntry entry)
        {
            if (entry.Directory.Type != ImageDebugType.CodeView)
                return entry;
 
            var reader = new BinaryReader(new MemoryStream(entry.Data));
            var newDataStream = new MemoryStream();
            var writer = new BinaryWriter(newDataStream);
 
            var sig = reader.ReadUInt32();
            if (sig != CodeViewSignature)
                return entry;
 
            writer.Write(sig);
            writer.Write(reader.ReadBytes(16)); // MVID
            writer.Write(reader.ReadUInt32()); // Age
 
            writer.Write(Encoding.UTF8.GetBytes(GetOriginalPdbPath()));
            writer.Write((byte)0);
 
            var newData = newDataStream.ToArray();
 
            var directory = entry.Directory;
            directory.SizeOfData = newData.Length;
 
            return new ImageDebugHeaderEntry(directory, newData);
        }
 
        string GetOriginalPdbPath()
        {
            if (!_module.HasDebugHeader)
                return string.Empty;
 
            var debugHeader = _module.GetDebugHeader();
            foreach (var entry in debugHeader.Entries)
            {
                if (entry.Directory.Type != ImageDebugType.CodeView)
                    continue;
 
                var reader = new BinaryReader(new MemoryStream(entry.Data));
                var sig = reader.ReadUInt32();
                if (sig != CodeViewSignature)
                    return string.Empty;
 
                var stream = reader.BaseStream;
                stream.Seek(16 + 4, SeekOrigin.Current); // MVID and Age
                                                         // Pdb path is NUL-terminated path at offset 24.
                                                         // https://github.com/dotnet/runtime/blob/main/docs/design/specs/PE-COFF.md#codeview-debug-directory-entry-type-2
                return Encoding.UTF8.GetString(
                    reader.ReadBytes((int)(stream.Length - stream.Position - 1))); // remaining length - ending \0
            }
 
            return string.Empty;
        }
 
        public ISymbolReaderProvider GetReaderProvider() => _symbolWriter.GetReaderProvider();
 
        public void Write(MethodDebugInformation info) => _symbolWriter.Write(info);
 
        public void Write(ICustomDebugInformationProvider provider) => _symbolWriter.Write(provider);
 
        public void Write() => _symbolWriter.Write();
 
        public void Dispose() => _symbolWriter.Dispose();
    }
}