using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Net.NetworkInformation;
using System.Runtime.InteropServices;
using MemScanEDR.Models;

namespace MemScanEDR.Services;

public class ProcessCollector
{
    private static readonly string[] SystemPaths = {
        @"C:\Windows\System32", @"C:\Windows\SysWOW64", @"C:\Windows"
    };

    private static readonly HashSet<string> SystemNames = new(StringComparer.OrdinalIgnoreCase) {
        "svchost.exe", "csrss.exe", "winlogon.exe", "services.exe",
        "lsass.exe", "spoolsv.exe", "smss.exe", "wininit.exe",
        "dwm.exe", "taskhostw.exe", "sihost.exe", "runtimebroker.exe",
        "explorer.exe", "system", "registry"
    };

    public List<ProcessInfo> CollectAll()
    {
        var processes = new List<ProcessInfo>();
        var connections = GetConnectionsByPid();

        foreach (var proc in Process.GetProcesses())
        {
            try
            {
                var info = new ProcessInfo
                {
                    Pid = proc.Id,
                    Name = proc.ProcessName,
                    ThreadCount = proc.Threads.Count,
                    WorkingSetMB = proc.WorkingSet64 / (1024 * 1024)
                };

                try { info.ExePath = proc.MainModule?.FileName ?? ""; } catch { }
                try { info.Ppid = GetParentPid(proc.Id); } catch { }

                if (info.Ppid > 0)
                {
                    try { info.ParentName = Process.GetProcessById(info.Ppid).ProcessName; }
                    catch { info.ParentName = "unknown"; }
                }

                try { info.CmdLine = GetCommandLine(proc.Id); } catch { }

                var exeLower = info.ExePath.ToLowerInvariant();
                info.IsSystemPath = SystemPaths.Any(p => exeLower.StartsWith(p.ToLowerInvariant()));
                info.IsTempPath = exeLower.Contains("\\temp\\") || exeLower.Contains("\\tmp\\") ||
                                  exeLower.Contains("\\appdata\\local\\temp\\");
                info.IsSigned = info.IsSystemPath;
                info.Signer = info.IsSystemPath ? "Microsoft Windows" : "";

                if (connections.TryGetValue(proc.Id, out var conns))
                    info.Connections = conns;

                processes.Add(info);
            }
            catch { }
            finally { proc.Dispose(); }
        }
        return processes;
    }

    private static int GetParentPid(int pid)
    {
        var handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if (handle == IntPtr.Zero) return 0;
        try
        {
            var entry = new PROCESSENTRY32 { dwSize = (uint)Marshal.SizeOf<PROCESSENTRY32>() };
            if (Process32First(handle, ref entry))
            {
                do
                {
                    if (entry.th32ProcessID == (uint)pid)
                        return (int)entry.th32ParentProcessID;
                } while (Process32Next(handle, ref entry));
            }
            return 0;
        }
        finally { CloseHandle(handle); }
    }

    private static string GetCommandLine(int pid)
    {
        try
        {
            using var searcher = new System.Management.ManagementObjectSearcher(
                $"SELECT CommandLine FROM Win32_Process WHERE ProcessId = {pid}");
            foreach (System.Management.ManagementObject obj in searcher.Get())
                return obj["CommandLine"]?.ToString() ?? "";
        }
        catch { }
        return "";
    }

    private static Dictionary<int, List<NetworkConnection>> GetConnectionsByPid()
    {
        var result = new Dictionary<int, List<NetworkConnection>>();
        try
        {
            var properties = IPGlobalProperties.GetIPGlobalProperties();
            foreach (var conn in properties.GetActiveTcpConnections())
            {
                var pid = GetConnectionPid(conn);
                if (pid > 0)
                {
                    if (!result.ContainsKey(pid)) result[pid] = new();
                    result[pid].Add(new NetworkConnection
                    {
                        Protocol = "tcp",
                        LocalAddr = conn.LocalEndPoint.ToString(),
                        RemoteAddr = conn.RemoteEndPoint.ToString(),
                        State = conn.State.ToString()
                    });
                }
            }
        }
        catch { }
        return result;
    }

    private static int GetConnectionPid(TcpConnectionInformation conn) => 0; // simplified

    private const uint TH32CS_SNAPPROCESS = 0x00000002;

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    private struct PROCESSENTRY32
    {
        public uint dwSize;
        public uint cntUsage;
        public uint th32ProcessID;
        public IntPtr th32DefaultHeapID;
        public uint th32ModuleID;
        public uint cntThreads;
        public uint th32ParentProcessID;
        public int pcPriClassBase;
        public uint dwFlags;
        [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 260)]
        public string szExeFile;
    }

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern IntPtr CreateToolhelp32Snapshot(uint dwFlags, uint th32ProcessID);

    [DllImport("kernel32.dll")]
    private static extern bool Process32First(IntPtr hSnapshot, ref PROCESSENTRY32 lppe);

    [DllImport("kernel32.dll")]
    private static extern bool Process32Next(IntPtr hSnapshot, ref PROCESSENTRY32 lppe);

    [DllImport("kernel32.dll")]
    private static extern bool CloseHandle(IntPtr hObject);
}