Main.cs

using System;
using System.Collections.Generic;
using System.Threading;
using Spectre.Console;

namespace IoMonitor
{
    /// <summary>
    /// Application entry point and top-level lifecycle management.
    /// </summary>
    internal static partial class Program
    {
        /// <summary>
        /// Entry point. Sets up security, ETW, sampling threads and runs the Spectre console UI.
        /// </summary>
        private static void Main()
        {
            bool shouldCleanup = false;

            try
            {
                if (!OperatingSystem.IsWindows())
                {
                    Console.WriteLine("DiskMonitor requires Windows because it reads Win32 process I/O counters and ETW events.");
                    return;
                }

                if (!AcquireSingleInstanceMutex())
                {
                    Console.WriteLine("I/O Monitor is already running. Only one instance can run at a time.");
                    return;
                }

                shouldCleanup = true;
                InitializeSecurityContext();
                AttachCancelKeyHandler();

                AnsiConsole.Clear();

                if (_hasAdmin)
                {
                    StartEtwMonitoring();
                    Thread.Sleep(1500);
                }

                StartDataCollectionThread();
                RunUiLoop();
            }
            catch (Exception ex)
            {
                AnsiConsole.WriteException(ex);
                WaitForKeyIfInteractive();
            }
            finally
            {
                if (shouldCleanup)
                {
                    Cleanup();
                }
            }
        }

        /// <summary>
        /// Attempts to acquire the single-instance mutex. Returns false if another instance is already running.
        /// </summary>
        private static bool AcquireSingleInstanceMutex()
        {
            bool createdNew;
            _singleInstanceMutex = new Mutex(initiallyOwned: true, name: "Global\\IoMonitorSingleton", out createdNew);
            return createdNew;
        }

        /// <summary>
        /// Populates fields that describe the current security context.
        /// </summary>
        private static void InitializeSecurityContext()
        {
            _hasAdmin = SecurityUtilities.IsAdministrator();
            _hasDebug = _hasAdmin && SecurityUtilities.EnableDebugPrivilege();
            _programStartTime = DateTime.Now;
        }

        /// <summary>
        /// Installs a Ctrl+C handler that triggers a graceful shutdown.
        /// </summary>
        private static void AttachCancelKeyHandler()
        {
            Console.CancelKeyPress += (_, e) =>
            {
                e.Cancel = true;
                _running = false;
            };
        }

        /// <summary>
        /// Starts the background sampling thread.
        /// </summary>
        private static void StartDataCollectionThread()
        {
            var dataThread = new Thread(DataCollectionLoop)
            {
                IsBackground = true,
                Name = "IoMonitor.DataCollection"
            };

            dataThread.Start();
        }

        /// <summary>
        /// Main UI loop that drives the live Spectre.Console layout.
        /// </summary>
        private static void RunUiLoop()
        {
            var layout = new Layout("Root").SplitRows(
                new Layout("Header").Size(1),
                new Layout("Status").Size(4),
                new Layout("Table"));

            AnsiConsole.Live(layout).AutoClear(false).Start(ctx =>
            {
                while (_running)
                {
                    HandleInput();

                    List<ProcessIoSample> snapshot;
                    lock (_dataLock)
                    {
                        snapshot = new List<ProcessIoSample>(_cachedSamples);
                    }

                    layout["Header"].Update(CreateHeaderPanel());
                    layout["Status"].Update(CreateStatusPanel(snapshot));
                    layout["Table"].Update(CreateTable(snapshot));
                    ctx.Refresh();

                    Thread.Sleep(250);
                }
            });
        }

        /// <summary>
        /// Performs best-effort cleanup of ETW resources, mutexes, and in-memory collections.
        /// </summary>
        private static void Cleanup()
        {
            AnsiConsole.MarkupLine("[yellow]Cleaning up resources...[/]");

            _running = false;

            StopEtwSession();
            FlushPendingCsvLog();

            try
            {
                _singleInstanceMutex?.ReleaseMutex();
                _singleInstanceMutex?.Dispose();
            }
            catch (Exception ex)
            {
                AnsiConsole.MarkupLine($"[red]Mutex cleanup error: {ex.Message}[/]");
            }

            try
            {
                lock (_dataLock)
                {
                    _history.Clear();
                    _etwTotals.Clear();
                    _etwLast.Clear();
                    _processIoByDrive.Clear();
                    _ioHistory.Clear();
                    _fileIoStats.Clear();
                    _cachedSamples.Clear();
                }
            }
            catch (Exception ex)
            {
                AnsiConsole.MarkupLine($"[red]Collection cleanup error: {ex.Message}[/]");
            }

            AnsiConsole.MarkupLine("[green]Cleanup complete. Exiting...[/]");
            Thread.Sleep(500);
        }

        /// <summary>
        /// Forces a synchronous resample using a nominal one-second interval.
        /// </summary>
        private static void ForceResample()
        {
            try
            {
                var samples = SampleProcesses(1.0, _showClosedProcesses);
                lock (_dataLock)
                {
                    _cachedSamples = samples;
                    _lastSampleIntervalSeconds = 1.0;
                }
            }
            catch
            {
                // Resample failures are ignored; the background loop will continue.
            }
        }

        private static void WaitForKeyIfInteractive()
        {
            try
            {
                if (!Console.IsInputRedirected)
                {
                    Console.ReadKey(intercept: true);
                }
            }
            catch
            {
                // Nothing useful to do when no console is attached.
            }
        }

        private static void FlushPendingCsvLog()
        {
            try
            {
                if (_csvLoggingOn)
                {
                    CsvExporter.AppendLogEntries("io_monitor.csv", DrainLogEntries());
                }
            }
            catch
            {
                // Shutdown logging is best effort.
            }
        }
    }
}