// 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 ();
	}
}