using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using MemScanEDR.Models;

namespace MemScanEDR.Services;

public class MemoryScanner
{
    public List<Signature> LoadSignatures(string rulesDir)
    {
        var sigs = GetBuiltinSignatures();
        // TODO: Load YARA rules from rulesDir
        return sigs;
    }

    public List<SignatureMatch> ScanProcess(int pid, List<Signature> signatures)
    {
        var matches = new List<SignatureMatch>();
        var hProcess = OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, false, pid);
        if (hProcess == IntPtr.Zero) return matches;

        try
        {
            foreach (var sig in signatures)
            {
                var results = ScanMemory(hProcess, sig);
                matches.AddRange(results);
            }
        }
        finally { CloseHandle(hProcess); }
        return matches;
    }

    private List<SignatureMatch> ScanMemory(IntPtr hProcess, Signature sig)
    {
        var results = new List<SignatureMatch>();
        var address = 0UL;
        var mbi = new MEMORY_BASIC_INFORMATION();
        var mbiSize = (uint)Marshal.SizeOf<MEMORY_BASIC_INFORMATION>();

        while (VirtualQueryEx(hProcess, (IntPtr)address, out mbi, mbiSize) != 0)
        {
            if (mbi.State == MEM_COMMIT && mbi.Protect != PAGE_NOACCESS && mbi.Protect != PAGE_GUARD)
            {
                var regionSize = (ulong)mbi.RegionSize;
                if (regionSize > 0 && regionSize < 100 * 1024 * 1024) // max 100MB per region
                {
                    var protect = GetProtectString(mbi.Protect);
                    var type = GetTypeString(mbi.Type);

                    // Check for suspicious patterns
                    if (sig.Category == "shellcode" || sig.Category == "c2" || sig.Category == "malware")
                    {
                        if (protect.Contains("EXECUTE") && type == "PRIVATE")
                        {
                            results.Add(new SignatureMatch
                            {
                                SigId = sig.SigId,
                                Category = sig.Category,
                                Severity = sig.Severity,
                                Description = sig.Description,
                                Address = $"0x{address:X}",
                                Protect = protect
                            });
                        }
                    }
                }
            }
            address = (ulong)mbi.BaseAddress + (ulong)mbi.RegionSize;
            if (address <= (ulong)mbi.BaseAddress) break; // overflow protection
        }
        return results;
    }

    public List<MemoryRegion> GetMemoryRegions(int pid)
    {
        var regions = new List<MemoryRegion>();
        var hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, false, pid);
        if (hProcess == IntPtr.Zero) return regions;

        try
        {
            var address = 0UL;
            var mbi = new MEMORY_BASIC_INFORMATION();
            var mbiSize = (uint)Marshal.SizeOf<MEMORY_BASIC_INFORMATION>();

            while (VirtualQueryEx(hProcess, (IntPtr)address, out mbi, mbiSize) != 0)
            {
                if (mbi.State == MEM_COMMIT)
                {
                    regions.Add(new MemoryRegion
                    {
                        BaseAddress = (ulong)mbi.BaseAddress,
                        Size = (ulong)mbi.RegionSize,
                        Protect = GetProtectString(mbi.Protect),
                        Type = GetTypeString(mbi.Type),
                        State = "COMMIT"
                    });
                }
                address = (ulong)mbi.BaseAddress + (ulong)mbi.RegionSize;
                if (address <= (ulong)mbi.BaseAddress) break;
            }
        }
        finally { CloseHandle(hProcess); }
        return regions;
    }

    public List<MemoryRegion> GetSuspiciousRegions(int pid)
    {
        var regions = GetMemoryRegions(pid);
        return regions.Where(r =>
            (r.Protect.Contains("EXECUTE_READWRITE") && r.Type == "PRIVATE") ||
            (r.Protect.Contains("EXECUTE") && r.Type != "IMAGE")
        ).ToList();
    }

    private static string GetProtectString(uint protect) => protect switch
    {
        0x02 => "READONLY",
        0x04 => "READWRITE",
        0x08 => "WRITECOPY",
        0x10 => "EXECUTE",
        0x20 => "EXECUTE_READ",
        0x40 => "EXECUTE_READWRITE",
        0x80 => "EXECUTE_WRITECOPY",
        _ => $"0x{protect:X2}"
    };

    private static string GetTypeString(uint type) => type switch
    {
        0x20000 => "PRIVATE",
        0x40000 => "MAPPED",
        0x1000000 => "IMAGE",
        _ => $"0x{type:X}"
    };

    public static List<Signature> GetBuiltinSignatures() => new()
    {
        new() { SigId = "SHELLCODE_NOP_SLED", Value = "\x90\x90\x90\x90", DataType = "string", Category = "shellcode", Severity = "medium", Description = "NOP sled detected" },
        new() { SigId = "SHELLCODE_GETEIP", Value = "\xE8\x00\x00\x00\x00\x58", DataType = "string", Category = "shellcode", Severity = "medium", Description = "CALL/POP EIP technique" },
        new() { SigId = "SHELLCODE_ROR13", Value = "\xC1\xC8\x0D", DataType = "string", Category = "shellcode", Severity = "medium", Description = "ROR13 API hash pattern" },
        new() { SigId = "C2_COBALTSTRIKE", Value = "beacon", DataType = "string", Category = "c2", Severity = "high", Description = "CobaltStrike Beacon string" },
        new() { SigId = "C2_METERPRETER", Value = "WS2_32.dll", DataType = "string", Category = "c2", Severity = "medium", Description = "Meterpreter WinSock" },
        new() { SigId = "INJECTION_REFLECTIVE", Value = "ReflectiveLoader", DataType = "string", Category = "injection", Severity = "high", Description = "Reflective DLL injection" },
        new() { SigId = "INJECTION_CRT", Value = "CreateRemoteThread", DataType = "string", Category = "injection", Severity = "medium", Description = "Remote thread injection" },
        new() { SigId = "INJECTION_VAE", Value = "VirtualAllocEx", DataType = "string", Category = "injection", Severity = "medium", Description = "VirtualAllocEx injection" },
        new() { SigId = "MALWARE_MIMIKATZ", Value = "mimikatz", DataType = "string", Category = "malware", Severity = "high", Description = "Mimikatz credential theft" },
        new() { SigId = "CRYPTO_MINER", Value = "stratum+tcp://", DataType = "string", Category = "crypto", Severity = "high", Description = "Crypto mining pool" },
        new() { SigId = "RANSOM_NOTE", Value = "YOUR_FILES_ARE_ENCRYPTED", DataType = "string", Category = "malware", Severity = "high", Description = "Ransomware note" },
        new() { SigId = "PS_DOWNLOAD", Value = "IEX (New-Object", DataType = "string", Category = "malware", Severity = "medium", Description = "PowerShell download exec" },
    };

    private const uint PROCESS_VM_READ = 0x0010;
    private const uint PROCESS_QUERY_INFORMATION = 0x0400;
    private const uint MEM_COMMIT = 0x1000;
    private const uint PAGE_NOACCESS = 0x01;
    private const uint PAGE_GUARD = 0x100;

    [StructLayout(LayoutKind.Sequential)]
    private struct MEMORY_BASIC_INFORMATION
    {
        public IntPtr BaseAddress;
        public IntPtr AllocationBase;
        public uint AllocationProtect;
        public IntPtr RegionSize;
        public uint State;
        public uint Protect;
        public uint Type;
    }

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern IntPtr OpenProcess(uint dwDesiredAccess, bool bInheritHandle, int dwProcessId);

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern uint VirtualQueryEx(IntPtr hProcess, IntPtr lpAddress, out MEMORY_BASIC_INFORMATION lpBuffer, uint dwLength);

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