using Microsoft.Diagnostics.Symbols;
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 System;
using System.Diagnostics;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Diagnostics.NETCore.Client;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Diagnostics.Tracing;

namespace Tracing
{
	/// <summary>
	/// A configurable, real-time EventPipe trace monitor for cross-platform use.
	/// This class uses a fluent builder pattern to attach to a .NET process,
	/// configure EventPipe providers, and process events in real-time.
	/// </summary>
	public class EventPipeTraceMonitor : IDisposable
	{
		private readonly int _pid;
		private int _durationSec = 10;
		private ClrTraceEventParser.Keywords _clrKeywords = ClrTraceEventParser.Keywords.None;
		private bool _monitorCpu = false;
		private readonly List<EventPipeProvider> _customProviders = new();

private Action<TraceEvent> _userCallback;
        private EventPipeSession _session;
        private TraceEventSource _traceLogSource; /// CHANGED
		private Task _processingTask;


		private Timer _timer;

		public EventPipeTraceMonitor( int processId )
		{
			_pid = processId;
		}

		// ####################################################################
		// Builder Configuration Methods
		// ####################################################################

		public EventPipeTraceMonitor WithDuration( int seconds )
		{
			_durationSec = seconds;
			return this;
		}

		/// <summary>
		/// Adds a custom EventPipeProvider.
		/// </summary>
		public EventPipeTraceMonitor AddProvider( EventPipeProvider provider )
		{
			_customProviders.Add( provider );
			return this;
		}

		// ####################################################################
		// Event Selection Methods
		// ####################################################################

		public EventPipeTraceMonitor MonitorExceptions()
		{
			// EventPipe requires JIT/Loader/Stack keywords to resolve symbols
			_clrKeywords |= ClrTraceEventParser.Keywords.Exception |
											ClrTraceEventParser.Keywords.Stack |
											ClrTraceEventParser.Keywords.Jit |
											ClrTraceEventParser.Keywords.JittedMethodILToNativeMap |
											ClrTraceEventParser.Keywords.Loader;
			return this;
		}

		public EventPipeTraceMonitor MonitorModuleLoads()
		{
			_clrKeywords |= ClrTraceEventParser.Keywords.Loader |
											ClrTraceEventParser.Keywords.Stack |
											ClrTraceEventParser.Keywords.Jit |
											ClrTraceEventParser.Keywords.Threading |
											ClrTraceEventParser.Keywords.PerfTrack |
											ClrTraceEventParser.Keywords.JittedMethodILToNativeMap;
			return this;
		}

		public EventPipeTraceMonitor MonitorCpuSamples()
		{
			_monitorCpu = true;
			return this;
		}

		// ####################################################################
		// Core Execution Methods
		// ####################################################################

		public void Start( Action<TraceEvent> onEventCallback )
		{
			_userCallback = onEventCallback;
			var providers = new List<EventPipeProvider>( _customProviders );





			if( _clrKeywords != ClrTraceEventParser.Keywords.None )
			{
				// Add the main CLR provider
				providers.Add( new EventPipeProvider(
								"Microsoft-Windows-DotNETRuntime",
								EventLevel.Informational,
								(long)_clrKeywords ) );
			}

			if( _monitorCpu )
			{
				// Add the CPU sampler provider
				providers.Add( new EventPipeProvider(
								"Microsoft-DotNETCore-SampleProfiler",
								EventLevel.Informational ) );
			}






			if( !providers.Any() )
			{
				log.warn( "No trace providers configured. Monitor will not start." );
				return;
			}

			try
			{
				var client = new DiagnosticsClient( _pid );
				log.info( $"Starting EventPipe session on PID {_pid} for {_durationSec}s..." );

				// Start the session. requestRundown=true is critical
				// to get symbols for code JIT-compiled before the session started.
				_session = client.StartEventPipeSession( providers, requestRundown: true );

				// Set up the timer to stop the session
				_timer = new Timer( delegate
				{
					log.info( $"Monitoring duration ({_durationSec} sec) elapsed." );
					Stop();
				}, null, _durationSec * 1000, Timeout.Infinite );

				// Create the TraceLogSource from the session's event stream
				_traceLogSource = new EventPipeEventSource( _session.EventStream );
				// KEEP TraceLog.CreateFromTraceEventSession( _session );

				// Register callbacks based on requested keywords
				if( ( _clrKeywords & ClrTraceEventParser.Keywords.Exception ) != 0 )
					_traceLogSource.Clr.ExceptionStart += OnEvent;
				if( ( _clrKeywords & ClrTraceEventParser.Keywords.Loader ) != 0 )
				{
					_traceLogSource.Clr.LoaderModuleLoad += OnEvent;
					_traceLogSource.Clr.All += OnAllEvents;

				}
				//if( _monitorCpu )
				//	_traceLogSource.CpuSpeedMHz += OnEvent;

				// Start processing the stream on a background thread.
				// .Process() is a blocking call that runs until the stream ends.
				_processingTask = Task.Run( () =>
				{
					try
					{/* _traceLogSource.Pro;*/ }
					catch( Exception ex ) { log.error( $"Trace processing error: {ex.Message}" ); }
					log.info( "Event processing finished." );
				} );
			}
			catch( Exception ex )
			{
				log.error( $"Failed to start EventPipe session: {ex.Message}" );
			}




		}

		private void OnAllEvents( TraceEvent evt )
		{
			log.info( $"EVENT: [PID:{evt.ProcessID}] -- {evt.EventName}" );
		}

		public void Stop()
		{
			if( _session != null )
			{
				log.info( "Stopping session..." );
				_timer?.Dispose();
				_timer = null;

				// Stopping the session will end the stream,
				// which causes traceLogSource.Process() to return.
				_session.Stop();
				_session.Dispose();
				_session = null;

				_traceLogSource?.Dispose();
				_traceLogSource = null;

				// Wait for the processing task to complete
				_processingTask?.Wait( 2000 );
				_processingTask = null;
			}
		}

		public void Dispose()
		{
			Stop();
		}

		/// <summary>
		/// Internal event handler.
		/// </summary>
		private void OnEvent( TraceEvent data )
		{
			// --- Pre-filtering ---
			if( data.Opcode == TraceEventOpcode.DataCollectionStart )
				return;
			if( data is ExceptionTraceData exd && exd.ExceptionType.Length == 0 )
				return;

			// --- Pass to user ---
			_userCallback( data );
		}
	}

	/// <summary>
	/// A configurable, real-time ETW trace monitor. (Windows-Only)
	/// This class uses a fluent builder pattern to configure and run a real-time
	/// TraceEventSession, processing events through a user-provided callback.
	/// </summary>
	public class RealTimeTraceMonitor : IDisposable
	{
		private TraceEventSession _session;
		private int _durationSec = 10;
		private string _processFilter = null;
		private KernelTraceEventParser.Keywords _kernelKeywords = KernelTraceEventParser.Keywords.ImageLoad | KernelTraceEventParser.Keywords.Process;
		private ClrTraceEventParser.Keywords _clrKeywords = ClrTraceEventParser.Keywords.None;
		private KernelTraceEventParser.Keywords _kernelStackKeywords = KernelTraceEventParser.Keywords.None;
		private bool _enableRundown = false;
		private bool _resolveNativeSymbols = false;
		private readonly List<Action<TraceLogEventSource>> _eventRegistrars = new();
		private Action<TraceEvent> _userCallback;
		private SymbolReader _symbolReader;
		private readonly object _symbolReaderLock = new();
		private TextWriter _symbolLog = new StringWriter();
		private Timer _timer;

		// ####################################################################
		// Builder Configuration Methods
		// ####################################################################

		/// <summary>
		/// Sets the monitoring duration in seconds.
		/// </summary>
		public RealTimeTraceMonitor WithDuration( int seconds )
		{
			_durationSec = seconds;
			return this;
		}

		/// <summary>
		/// Filters events to only include those from a process whose name contains this string.
		/// </summary>
		public RealTimeTraceMonitor FilterByProcess( string processName )
		{
			_processFilter = processName;
			return this;
		}

		/// <summary>
		/// Provides a TextWriter for detailed symbol lookup messages.
		/// </summary>
		public RealTimeTraceMonitor WithSymbolLog( TextWriter logWriter )
		{
			_symbolLog = logWriter;
			return this;
		}

		// ####################################################################
		// Event Selection Methods
		// ####################################################################

		/// <summary>
		/// Enables monitoring for CLR Exceptions (ExceptionStart).
		/// Automatically enables JIT, Loader, and Stack keywords required for symbol resolution.
		/// </summary>
		public RealTimeTraceMonitor MonitorExceptions()
		{
			_clrKeywords |= ClrTraceEventParser.Keywords.Exception |
											ClrTraceEventParser.Keywords.Stack |
											ClrTraceEventParser.Keywords.Jit |
											ClrTraceEventParser.Keywords.JittedMethodILToNativeMap |
											ClrTraceEventParser.Keywords.Loader;
			_enableRundown = true;
			_resolveNativeSymbols = true; // Native symbol resolution is often needed for full exception stacks
			_eventRegistrars.Add( source => source.Clr.ExceptionStart += OnEvent );
			return this;
		}

		/// <summary>
		/// Enables monitoring for CLR Module Loads (LoaderModuleLoad).
		/// Automatically enables JIT, Loader, and Stack keywords.
		/// </summary>
		public RealTimeTraceMonitor MonitorModuleLoads()
		{
			_clrKeywords |= ClrTraceEventParser.Keywords.Loader |
											ClrTraceEventParser.Keywords.Stack |
											ClrTraceEventParser.Keywords.Jit |
											ClrTraceEventParser.Keywords.JittedMethodILToNativeMap;
			_enableRundown = true;
			_eventRegistrars.Add( source => source.Clr.LoaderModuleLoad += OnEvent );
			return this;
		}

		/// <summary>
		/// Enables monitoring for CPU Samples (PerfInfoSample).
		/// </summary>
		public RealTimeTraceMonitor MonitorCpuSamples()
		{
			_kernelKeywords |= KernelTraceEventParser.Keywords.Profile;
			_kernelStackKeywords |= KernelTraceEventParser.Keywords.Profile;
			_resolveNativeSymbols = true; // CPU samples require native symbol resolution
			_eventRegistrars.Add( source => source.Kernel.PerfInfoSample += OnEvent );
			return this;
		}

		/// <summary>
		/// Enables monitoring for arbitrary Kernel events.
		/// </summary>
		/// <param name="events">Kernel keywords to enable.</param>
		/// <param name="stackEvents">Keywords to collect stacks for.</param>
		public RealTimeTraceMonitor MonitorKernelEvents( KernelTraceEventParser.Keywords events, KernelTraceEventParser.Keywords stackEvents = KernelTraceEventParser.Keywords.None )
		{
			_kernelKeywords |= events;
			_kernelStackKeywords |= stackEvents;
			if( stackEvents != KernelTraceEventParser.Keywords.None )
			{
				_resolveNativeSymbols = true;
			}
			return this;
		}

		/// <summary>
		/// Enables monitoring for arbitrary CLR events.
		/// Note: For complex events, prefer the specific Monitor* methods.
		/// </summary>
		/// <param name="events">CLR keywords to enable.</param>
		public RealTimeTraceMonitor MonitorClrEvents( ClrTraceEventParser.Keywords events )
		{
			_clrKeywords |= events;
			if( ( events & ClrTraceEventParser.Keywords.Stack ) != 0 )
			{
				_clrKeywords |= ClrTraceEventParser.Keywords.Jit |
												ClrTraceEventParser.Keywords.JittedMethodILToNativeMap |
												ClrTraceEventParser.Keywords.Loader;
				_enableRundown = true;
				_resolveNativeSymbols = true;
			}
			return this;
		}

		// ####################################################################
		// Core Execution Methods
		// ####################################################################

		/// <summary>
		/// Starts the monitoring session and blocks until the duration expires or Ctrl+C is pressed.
		/// </summary>
		/// <param name="onEventCallback">The callback to execute for each filtered event.</param>
		public void Start( Action<TraceEvent> onEventCallback )
		{
			if( !OperatingSystem.IsWindows() )
			{
				log.error( $"{nameof( RealTimeTraceMonitor )} is only supported on Windows (ETW)." );
				return;
			}

			_userCallback = onEventCallback;
			log.info( $"Starting {nameof( RealTimeTraceMonitor )} for {_durationSec} seconds." );

			if( TraceEventSession.GetActiveSessionNames().Contains( "TraceLogMonitorSession" ) )
			{
				log.warn( "Session 'TraceLogMonitorSession' is already active. Attaching." );
			}

			try
			{
				using( _session = new TraceEventSession( "TraceLogMonitorSession", TraceEventSessionOptions.Attach ) )
				{
					// Set up Ctrl-C to stop the session
					Console.CancelKeyPress += ( object sender, ConsoleCancelEventArgs cancelArgs ) =>
					{
						cancelArgs.Cancel = true;
						Stop();
					};

					SetupProviders();
					SetupSymbolReader();

					using( var traceLogSource = TraceLog.CreateFromTraceEventSession( _session ) )
					{
						// Register all requested event callbacks
						foreach( var registrar in _eventRegistrars )
						{
							registrar( traceLogSource );
						}

						// Set up a timer to stop processing after the duration
						_timer = new Timer( delegate ( object state )
						{
							log.info( $"Monitoring duration ({_durationSec} sec) elapsed." );
							Stop();
						}, null, _durationSec * 1000, Timeout.Infinite );

						log.info( "Processing events..." );
						traceLogSource.Process();
						log.info( "Event processing finished." );
					}
				}
			}
			catch( Exception ex )
			{
				log.error( $"Failed to start trace session: {ex.Message}" );
				log.error( "This often requires Admin privileges." );
			}
			finally
			{
				_session = null;
			}
		}

		/// <summary>
		/// Stops the monitoring session.
		/// </summary>
		public void Stop()
		{
			if( _session != null )
			{
				log.info( "Stopping session..." );
				_session.Dispose();
				_session = null;
			}
			if( _timer != null )
			{
				_timer.Dispose();
				_timer = null;
			}
		}

		public void Dispose()
		{
			Stop();
		}

		// ####################################################################
		// Private Helper Methods
		// ####################################################################

		private void SetupProviders()
		{
			log.info( $"Enabling Kernel Providers with keywords: {_kernelKeywords}" );
			try
			{
				_session.EnableKernelProvider( _kernelKeywords, _kernelStackKeywords );
			}
			catch( Exception ex )
			{
				log.error( $"Failed to enable Kernel provider: {ex}" );
			}

			if( _clrKeywords != ClrTraceEventParser.Keywords.None )
			{
				log.info( $"Enabling CLR Provider with keywords: {_clrKeywords}" );
				try
				{
					_session.EnableProvider(
							ClrTraceEventParser.ProviderGuid,
							TraceEventLevel.Informational,
							(ulong)_clrKeywords );
				}
				catch( Exception ex )
				{
					log.error( $"Failed to enable CLR provider: {ex}" );
				}
			}

			if( _enableRundown )
			{
				// Remove keywords not relevant for rundown
				var rundownKeywords = ( _clrKeywords & ~( ClrTraceEventParser.Keywords.Exception | ClrTraceEventParser.Keywords.Stack ) );
				rundownKeywords |= ClrTraceEventParser.Keywords.StartEnumeration;

				log.info( $"Enabling CLR Rundown Provider with keywords: {rundownKeywords}" );
				try
				{
					_session.EnableProvider(
							ClrRundownTraceEventParser.ProviderGuid,
							TraceEventLevel.Informational,
							(ulong)rundownKeywords );
				}
				catch( Exception ex )
				{
					log.error( $"Failed to enable CLR Rundown provider: {ex}" );
				}
			}
		}

		private void SetupSymbolReader()
		{
			if( _resolveNativeSymbols )
			{
				log.info( "Setting up SymbolReader for native symbol resolution." );
				var symbolPath = new SymbolPath( SymbolPath.SymbolPathFromEnvironment ).Add( SymbolPath.MicrosoftSymbolServerPath );
				_symbolReader = new SymbolReader( _symbolLog, symbolPath.ToString() );
				// Allow reading PDBs from "unsafe" locations (next to EXE)
				_symbolReader.SecurityCheck = ( path => true );
			}
		}

		/// <summary>
		/// Primary event handler.
		/// This method is called from multiple threads and must be thread-safe.
		/// </summary>
		private void OnEvent( TraceEvent data )
		{
			// --- Pre-filtering ---
			if( data.Opcode == TraceEventOpcode.DataCollectionStart )
				return;
			if( data is ExceptionTraceData exd && exd.ExceptionType.Length == 0 )
				return;

			// --- Process filter ---
			if( _processFilter != null && !data.ProcessName.Contains( _processFilter, StringComparison.OrdinalIgnoreCase ) )
			{
				return;
			}

			// --- Symbol Resolution ---
			var callStack = data.CallStack();
			if( callStack != null && _resolveNativeSymbols )
			{
				ResolveNativeStack( callStack );
			}

			// --- Pass to user ---
			_userCallback( data );
		}

		/// <summary>
		/// Resolves native symbols for a given call stack.
		/// This method locks the SymbolReader, as it is not thread-safe.
		/// </summary>
		private void ResolveNativeStack( TraceCallStack callStack )
		{
			if( _symbolReader == null )
				return;

			lock( _symbolReaderLock )
			{
				while( callStack != null )
				{
					var codeAddress = callStack.CodeAddress;
					if( codeAddress.Method == null )
					{
						var moduleFile = codeAddress.ModuleFile;
						if( moduleFile != null )
						{
							codeAddress.CodeAddresses.LookupSymbolsForModule( _symbolReader, moduleFile );
						}
					}
					callStack = callStack.Caller;
				}
			}
		}
	}

	/// <summary>
	/// Example implementation of the RealTimeTraceMonitor utility.
	/// </summary>
	internal class TraceLogMonitor
	{
		public static void Run()
		{
			log.info( "******************** RealTimeTraceLog DEMO ********************" );

			// Define a thread-safe callback to process events
			Action<TraceEvent> printCallback = ( TraceEvent data ) =>
			{
				// log.info is assumed to be thread-safe
				// Note: EventPipe data does not have ProcessName, as it's scoped to the process.
				log.info( $"EVENT: [PID:{data.ProcessID}] -- {data.EventName}" );

				var stack = data.CallStack();
				if( stack != null )
				{
					log.info( $"CALLSTACK: {stack}" );
				}
			};

			// Run an exception generator in the background
			Task.Factory.StartNew( delegate
			{
				Thread.Sleep( 3000 );
				ThrowException();
			} );

			try
			{
				if( OperatingSystem.IsWindows() )
				{
					log.info( "Windows detected. Using ETW-based RealTimeTraceMonitor." );
					using( var monitor = new RealTimeTraceMonitor() )
					{
						monitor.WithDuration( 10 )
									 .MonitorExceptions()
									 .MonitorModuleLoads()
									 .Start( printCallback );
					}
				}
				else
				{
					log.info( "Non-Windows detected. Using EventPipe-based EventPipeTraceMonitor." );
					int currentPid = Process.GetCurrentProcess().Id;
					using( var monitor = new EventPipeTraceMonitor( currentPid ) )
					{
						monitor.WithDuration( 10 )
									 .MonitorExceptions()
									 .MonitorModuleLoads()
									 //.MonitorCpuSamples() // Note: EventPipe CPU sampling is very noisy
									 .Start( printCallback );
					}
				}
			}
			catch( Exception ex )
			{
				log.error( $"Monitoring failed: {ex}" );
			}

			log.info( "Finished monitoring." );
		}

		// --- Exception generation helpers ---

		[System.Runtime.CompilerServices.MethodImpl( System.Runtime.CompilerServices.MethodImplOptions.NoInlining )]
		private static void ThrowException()
		{
			ThrowException1();
		}

		[System.Runtime.CompilerServices.MethodImpl( System.Runtime.CompilerServices.MethodImplOptions.NoInlining )]
		private static void ThrowException1()
		{
			log.info( "Causing an exception to happen so a CLR Exception Start event will be generated." );
			try
			{
				throw new Exception( "This is a test exception thrown to generate a CLR event" );
			}
			catch( Exception ) { }
		}
	}
}