using Microsoft.Diagnostics.NETCore.Client;
using Microsoft.Diagnostics.Tracing;
using Microsoft.Diagnostics.Tracing.Etlx;
using Microsoft.Diagnostics.Tracing.Parsers;
using Microsoft.Diagnostics.Tracing.Parsers.Clr;
using Microsoft.Diagnostics.Tracing.Parsers.Kernel;
using Microsoft.Diagnostics.Tracing.Session;
using ProfilerHelpers;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.Tracing;
using System.Threading;
using System.Threading.Tasks;
using TraceKeywords = Microsoft.Diagnostics.Tracing.Parsers.ClrTraceEventParser.Keywords;

namespace Tracing
{
	public class MemoryPipe
	{
		private readonly DiagnosticsClient _client;
		private readonly PerProcessProfilingState _processes;

		// because we are not interested in self monitoring
		private readonly int _currentPid;

		private int _started = 0;

		Thread _thread;

		public MemoryPipe()
		{

			_currentPid = Process.GetCurrentProcess().Id;

			_client = new DiagnosticsClient( _currentPid );

			_processes = new();


			var threadStart = new ThreadStart( () => Start( true ) );
			_thread = new Thread( threadStart );
			_thread.Start();


		}

		//public async Task StartAsync( bool allAllocations )
		public void Start( bool allAllocations )
		{
			if( Interlocked.CompareExchange( ref _started, 1, 0 ) == 1 )
			{
				throw new InvalidOperationException( "Impossible to start profiling more than once." );
			}

			log.info( $"Starting MemoryPipe on thread {Thread.CurrentThread.ManagedThreadId}" );

			var providers = new List<EventPipeProvider>()
				{
						new EventPipeProvider("Microsoft-Windows-DotNETRuntime",
								EventLevel.Verbose, (long)(
									TraceKeywords.GC |
									TraceKeywords.Contention |
									TraceKeywords.Debugger |
									TraceKeywords.Exception |
									TraceKeywords.GCAllObjectAllocation |
									TraceKeywords.GCSampledObjectAllocationHigh |
									TraceKeywords.GCSampledObjectAllocationLow |
									TraceKeywords.Security |
									TraceKeywords.Threading |
									TraceKeywords.Type |
									TraceKeywords.TypeDiagnostic |
									TraceKeywords.WaitHandle |
									TraceKeywords.All
									) )
				};


			//await Task.Factory.StartNew( () => {
				using EventPipeSession session = _client.StartEventPipeSession( providers, false );

					//log.info( $"SetupProviders" );
					//SetupProviders( session, allAllocations );
					var source = new EventPipeEventSource( session.EventStream );


					log.info( $"SetupListeners" );
					SetupListeners( source );

					log.info( $"Run Process" );
					source.Process();

					log.info( $"Done" );
			//} );
		}

		private void SetupProviders( TraceEventSession session, bool noSampling )
		{
			// Note: the kernel provider MUST be the first provider to be enabled
			// If the kernel provider is not enabled, the callstacks for CLR events are still received
			// but the symbols are not found (except for the application itself)
			// Maybe a TraceEvent implementation details triggered when a module (image) is loaded
			var success = true;

			//*
			log.info( $"EnableKernelProvider" );
			success = log.call( session.EnableKernelProvider( KernelTraceEventParser.Keywords.None |
//					KernelTraceEventParser.Keywords.ImageLoad |
//					KernelTraceEventParser.Keywords.Process
					0 ,
					KernelTraceEventParser.Keywords.None
			) );
			log.info( $"EnableKernelProvider {success}" );
			//*/

			// The CLR source code indicates that the provider must be set before the monitored application starts
			// Note: no real difference between High and Low
			ClrTraceEventParser.Keywords eventsKeyword = noSampling
					? ClrTraceEventParser.Keywords.GCSampledObjectAllocationLow | ClrTraceEventParser.Keywords.GCSampledObjectAllocationHigh
					: ClrTraceEventParser.Keywords.GCSampledObjectAllocationLow
					;

			log.info( $"EnableProvider" );
			success = log.call( session.EnableProvider(
					ClrTraceEventParser.ProviderGuid,
					TraceEventLevel.Verbose,    // this is needed in order to receive GCSampledObjectAllocation event
					(ulong)(

					eventsKeyword |

					// required to receive the BulkType events that allows
					// mapping between the type ID received in the allocation events
					ClrTraceEventParser.Keywords.GCHeapAndTypeNames |
					ClrTraceEventParser.Keywords.Type |

					// events related to JITed methods
					ClrTraceEventParser.Keywords.Jit |              // Turning on JIT events is necessary to resolve JIT compiled code
					ClrTraceEventParser.Keywords.JittedMethodILToNativeMap | // This is needed if you want line number information in the stacks
					ClrTraceEventParser.Keywords.Loader |                   // You must include loader events as well to resolve JIT compiled code.

					// this is mandatory to get the callstacks in each CLR event payload.
					//ClrTraceEventParser.Keywords.Stack |

					0
					)
			) );
			log.info( $"EnableProvider {success}" );


			// Note: ClrRundown is not needed because only new processes will be monitored
		}

		private void SetupListeners( EventPipeEventSource source )
		{
			// register for high and low keyword
			// if both are set, each allocation will trigger an event (beware performance issues...)
			//source.Clr.GCSampledObjectAllocation += OnSampleObjectAllocation;
			source.Clr.GCAllocationTick += OnAllocTick;

			// required to receive the mapping between type ID (received in GCSampledObjectAllocation)
			// and their name (received in TypeBulkType)
			source.Clr.TypeBulkType += OnTypeBulkType;

			// messages to get callstacks
			// the correlation seems to be as "simple" as taking the last event on the same thread
			source.Clr.ClrStackWalk += OnClrStackWalk;

			// needed to get JITed method details
			source.Clr.MethodLoadVerbose += OnMethodDetails;
			source.Clr.MethodDCStartVerboseV2 += OnMethodDetails;

			source.Clr.ContentionLockCreated += OnLockCreated;
			source.Clr.ContentionStart += OnLockStart;
			source.Clr.ContentionStop += OnLockStop;

			// get notified when a module is load to map the corresponding symbols
			source.Kernel.ImageLoad += OnImageLoad;
		}

		private void OnAllocTick( GCAllocationTickTraceData data )
		{
		if( FilterOutEvent( data ) )
				return;

			//log.info( $"*** RAW: {data}" );

			//var callStack = data.CallStack();
			//var caller = callStack.Caller;

			//log.info( $"Call stack {callStack}" );

			var thisThreadId = Thread.CurrentThread.ManagedThreadId;

			var typeName = GetProcessTypeName( data.ProcessID, data.TypeID );
			//if( data.TotalSizeForTypeSample >= 85000 )
			{
				var message = $"{data.ThreadID}/{thisThreadId} Alloc {data.ObjectSize, 8} at 0x{data.Address:0000000000000000} in {data.TypeName} (or {GetProcessTypeName( data.ProcessID, data.TypeID )}) ";
				log.info( message );
			}
			GetProcessAllocations( data.ProcessID )
					.AddAllocation(
					data.ThreadID,
					(ulong)data.ObjectSize,
					(ulong)1,
					typeName
					);
		}

		private void OnLockCreated( ContentionLockCreatedTraceData data )
		{
			log.info( $"{data}" );
		}

		private void OnLockStart( ContentionStartTraceData data )
		{
			log.info( $"{data}" );
		}

		private void OnLockStop( ContentionStopTraceData data )
		{
			log.info( $"{data}" );
		}

		private void OnImageLoad( ImageLoadTraceData data )
		{
			//if( FilterOutEvent( data ) )
			//	return;

			log.info( $"{data}" );

			//GetProcessMethods( data.ProcessID ).AddModule( data.FileName, data.ImageBase, data.ImageSize );

			log.info($"{data.ProcessID}.{data.ThreadID} --> {data.FileName}");
		}

		private void OnMethodDetails( MethodLoadUnloadVerboseTraceData data )
		{
			//if( FilterOutEvent( data ) )
		//		return;

			//log.info( $"{data}" );

			// care only about jitted methods
			if( !data.IsJitted )
				return;

			var method = GetProcessMethods( data.ProcessID )
					.Add( data.MethodStartAddress, data.MethodSize, data.MethodNamespace, data.MethodName, data.MethodSignature );

			//log.info( $"0x{data.MethodStartAddress.ToString( "x12" )} - {data.MethodSize,6} | {data.MethodName}" );
		}

		private MethodStore GetProcessMethods( int pid )
		{
			if( !_processes.Methods.TryGetValue( pid, out var methods ) )
			{
				methods = new MethodStore( pid );
				_processes.Methods[pid] = methods;
			}
			return methods;
		}


		private void OnSampleObjectAllocation( GCSampledObjectAllocationTraceData data )
		{
			if( FilterOutEvent( data ) )
				return;

			//log.info( $"{data}" );

			var typeName = GetProcessTypeName( data.ProcessID, data.TypeID );
			//if( data.TotalSizeForTypeSample >= 85000 )
			{
				var message = $"{data.ProcessID}.{data.ThreadID} - {data.TimeStampRelativeMSec,12} | Alloc {GetProcessTypeName( data.ProcessID, data.TypeID )} ({data.TotalSizeForTypeSample})";
				log.info( message );
			}
			GetProcessAllocations( data.ProcessID )
					.AddAllocation(
					data.ThreadID,
					(ulong)data.TotalSizeForTypeSample,
					(ulong)data.ObjectCountForTypeSample,
					typeName
					);
		}

		private ProcessAllocations GetProcessAllocations( int pid )
		{
			if( !_processes.Allocations.TryGetValue( pid, out var allocations ) )
			{
				allocations = new ProcessAllocations( pid );
				_processes.Allocations[pid] = allocations;
			}
			return allocations;
		}

		private void OnClrStackWalk( ClrStackWalkTraceData data )
		{
			var message = $"{data.ProcessID}.{data.ThreadID} - {data.TimeStampRelativeMSec,12} |       {data.FrameCount} frames";
			log.info(message);

			var callstack = BuildCallStack( data );
			GetProcessAllocations( data.ProcessID ).AddStack( data.ThreadID, callstack );
			//DumpStack(data);
		}

		private AddressStack BuildCallStack( ClrStackWalkTraceData data )
		{
			//log.info( $"{data}" );

			var length = data.FrameCount;
			AddressStack stack = new AddressStack( length );

			// frame 0 is the last frame of the stack (i.e. last called method)
			for( int i = 0; i < length; i++ )
			{
				stack.AddFrame( data.InstructionPointer( i ) );
			}

			return stack;
		}

		private void DumpStack( ClrStackWalkTraceData data )
		{
			if( FilterOutEvent( data ) )
				return;

			log.info( $"{data}" );

			var methods = GetProcessMethods( data.ProcessID );
			for( int i = 0; i < data.FrameCount; i++ )
			{
				var address = data.InstructionPointer( i );
				log.info( methods.GetFullName( address ) );
			}
			log.info( $"" );
		}

		private void OnTypeBulkType( GCBulkTypeTraceData data )
		{

			ProcessTypeMapping mapping = GetProcessTypesMapping( data.ProcessID );
			for( int currentType = 0; currentType < data.Count; currentType++ )
			{
				GCBulkTypeValues value = data.Values( currentType );
				mapping[value.TypeID] = value.TypeName;
				//log.info( $"{value}" );
			}
		}

		private ProcessTypeMapping GetProcessTypesMapping( int pid )
		{
			ProcessTypeMapping mapping;
			if( !_processes.Types.TryGetValue( pid, out mapping ) )
			{
				AssociateProcess( pid );

				mapping = new ProcessTypeMapping( pid );
				_processes.Types[pid] = mapping;
			}
			return mapping;
		}

		private void AssociateProcess( int pid )
		{
			try
			{
				_processes.Names[pid] = Process.GetProcessById( pid ).ProcessName;
			}
			catch( Exception )
			{
				log.info( $"? {pid}" );
				// we might not have access to the process
			}
		}

		private string GetProcessTypeName( int pid, ulong typeID )
		{
			if( !_processes.Types.TryGetValue( pid, out var mapping ) )
			{
				return typeID.ToString();
			}

			var name = mapping[typeID];
			return string.IsNullOrEmpty( name ) ? typeID.ToString() : name;
		}

		private bool FilterOutEvent( TraceEvent data )
		{
			return data.ThreadID == Thread.CurrentThread.ManagedThreadId;
			//return false; // ( data.ProcessID == _currentPid );
		}
	}
}