using System;
using System.Collections.Generic;
using System.IO;
using System.Management;
using System.Text;
using System.Threading;
using System.Threading.Tasks;


namespace InsiderAttackdetector
{
    public class DataGatherer : IDisposable
    {
        private StringBuilder log = new StringBuilder();

        private readonly string LOG_FILE_SAVE_PATH = Path.Combine(Config.LOGS_DIR_PATH, $"log-{Guid.NewGuid()}-{DateTime.UtcNow.ToLocalTime().ToString("yyyyMMdd--HH-mm-ss")}.csv");

        #region PROCESS_TRACES
        private ManagementEventWatcher processStartTrace;

        private readonly string processStartTraceQuery = "SELECT * FROM Win32_ProcessStartTrace";

        private ManagementEventWatcher processStopTrace;

        private readonly string processStopTraceQuery = "SELECT * FROM Win32_ProcessStopTrace";
        #endregion

        #region FILTERS
        private List<string> safeProcesses = new();
        
        private List<string> unwantedProperties = new() { "SessionID", "ThreadID", "ExitStatus"};
        #endregion


        public DataProcessor DataProcessor { get; set; }


        public DataGatherer()
        {
            DataProcessor = new DataProcessor();
            
            InitializeWindowsProcesses();

            InitializeWatchers();

            InitializeLogFile();
        }


        public async Task<string> Capture()
        {
            Task.Run(() => InitiateFlush());

            StartWatchers();

            Console.WriteLine("Press 'q' to exit data capturing!");

            while (Console.ReadKey().KeyChar != 'q') Thread.Sleep(5000);

            return LOG_FILE_SAVE_PATH;
        }


        #region INITIALIZERS
        private void InitializeWindowsProcesses()
        {
            var safeProcessesString = File.ReadAllLines(Config.SAFE_PROCESSES_PATH)[0];

            safeProcesses.AddRange(safeProcessesString.Split(","));
        }


        private void InitializeWatchers()
        {
            processStartTrace = GetEventWatcher(processStartTraceQuery, new EventArrivedEventHandler(ProcessStart_EventArrived));

            processStopTrace = GetEventWatcher(processStopTraceQuery, new EventArrivedEventHandler(ProcessStop_EventArrived));
        }


        private void InitializeLogFile()
        {
            Console.WriteLine(Config.CSV_HEADER);

            File.AppendAllText(LOG_FILE_SAVE_PATH, Config.CSV_HEADER);
        }
        #endregion


        #region WATCHER_METHODS
        private void StartWatchers()
        {
            processStartTrace.Start();

            processStopTrace.Start();
        }


        private void StopWatchers()
        {
            processStartTrace.Stop();

            processStopTrace.Stop();
        }


        private ManagementEventWatcher GetEventWatcher(string query, EventArrivedEventHandler eventHandler)
        {
            var eventWatcher = new ManagementEventWatcher(new WqlEventQuery(query));

            eventWatcher.EventArrived += eventHandler;

            return eventWatcher;
        }


        private void ProcessStop_EventArrived(object sender, EventArrivedEventArgs e)
        {
            Thread.Sleep(10); // In case of fast start and stop sequence of processes, let this action sleep for a brief moment, so we get proper order.

            if (ProcessMessageAllAttributes(e, out var message, true))
            {
                Console.WriteLine($"Process stopped: {message}");
            }
        }


        private void ProcessStart_EventArrived(object sender, EventArrivedEventArgs e)
        {
            if (ProcessMessageAllAttributes(e, out var message))
            {
                log.Append($"procStart,{message}{Environment.NewLine}");
                Console.WriteLine($"Process started: {message}");
            }
        }
        #endregion


        #region LOG_FLUSH_LOGIC
        private async Task FlushLog()
        {
            await File.AppendAllTextAsync(LOG_FILE_SAVE_PATH, log.ToString());

            log.Clear();
        }


        private async Task InitiateFlush()
        {
            while (true)
            {
                Console.WriteLine($"Saving log start: {DateTime.UtcNow}");

                await FlushLog();

                Console.WriteLine($"Saving log end: {DateTime.UtcNow}");

                Thread.Sleep(Config.FLUSH_PERIOD);
            }
        }
        #endregion


        private bool ProcessMessageAllAttributes(EventArrivedEventArgs e, out string message, bool isStop = false)
        {
            var messageBuilder = new StringBuilder();
            message = "";

            int parent = -1;
            int child = -1;

            foreach (var prop in e.NewEvent.Properties)
            {
                if (prop.Name == "ParentProcessID")
                {
                    parent = int.Parse(prop.Value.ToString());
                    continue;
                }

                if (prop.Name == "ProcessID")
                {
                    child = int.Parse(prop.Value.ToString());
                    continue;
                }

                if (prop.Name == "TIME_CREATED")
                {
                    // Timestamp
                    var date = Helpers.GetDateFromTicks(long.Parse(prop.Value.ToString()));
                    messageBuilder.Append($"{date},");
                    continue;
                }

                // Ignore properties
                if (unwantedProperties.Contains(prop?.Name))
                {
                    continue;
                }

                // Ignore safe processes
                if (safeProcesses.Contains(prop.Value?.ToString()))
                {
                    return false;
                }
                
                if (!prop.IsArray)
                {
                    messageBuilder.Append($"{prop.Value},");
                }
            }

            var parentChildMessage = DataProcessor.ProcessFamilyRelationShip(child, parent, isStop);
            if (!String.IsNullOrWhiteSpace(parentChildMessage))
            {
                messageBuilder.Append(parentChildMessage);
            }

            message = Helpers.TrimComma(messageBuilder.ToString());

            return true;
        }


        public void Dispose()
        {
            StopWatchers();
        }
    }
}
