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

namespace MemScanEDR.Services;

/// <summary>
/// 熵值分析器 —— 计算内存区域的 Shannon 熵
/// 高熵 (>7.0) 通常表示加密/压缩载荷，是恶意软件的重要指标
/// 正常程序代码段熵值通常稳定在 3.0-6.5 范围
/// </summary>
public class EntropyAnalyzer
{
    // Entropy calculation constants
    private const int SampleSize = 4096;       // 默认采样大小
    private const int MaxSamplesPerRegion = 8; // 每个区域最多采样点
    private const double HighEntropyThreshold = 7.0;

    private const uint PROCESS_VM_READ = 0x0010;

    /// <summary>
    /// 计算单个内存区域的熵值（通过采样子区域）
    /// </summary>
    public double CalculateRegionEntropy(IntPtr hProcess, ulong baseAddress, ulong regionSize)
    {
        if (regionSize == 0) return 0;

        var entropies = new List<double>();
        var sampleCount = Math.Min(MaxSamplesPerRegion, (int)(regionSize / SampleSize));
        if (sampleCount < 1) sampleCount = 1;

        for (int i = 0; i < sampleCount; i++)
        {
            var offset = (ulong)(i * (long)(regionSize / (ulong)sampleCount));
            var readSize = Math.Min(SampleSize, regionSize - offset);
            if (readSize > int.MaxValue) readSize = SampleSize;

            var buffer = new byte[readSize];
            if (ReadProcessMemory(hProcess, (IntPtr)(baseAddress + offset), buffer, (int)readSize, out _))
            {
                entropies.Add(CalculateShannonEntropy(buffer));
            }
        }

        return entropies.Count > 0 ? entropies.Average() : 0;
    }

    /// <summary>
    /// 批量计算所有内存区域的熵值，更新 MemoryRegion.Entropy 字段
    /// </summary>
    public EntropyResult AnalyzeRegions(IntPtr hProcess, List<MemoryRegion> regions)
    {
        var result = new EntropyResult();
        var allEntropies = new List<double>();

        foreach (var region in regions)
        {
            // 只对已提交的可读区域计算熵
            if (region.State != "COMMIT") continue;
            if (region.Protect == "NOACCESS" || region.Protect == "0x01") continue;

            region.Entropy = CalculateRegionEntropy(hProcess, region.BaseAddress, region.Size);
            allEntropies.Add(region.Entropy);

            if (region.Entropy > HighEntropyThreshold)
            {
                result.HighEntropyRegions++;
                result.HighEntropyTotalSize += region.Size;
            }
        }

        if (allEntropies.Count > 0)
        {
            result.AvgEntropy = allEntropies.Average();
            result.MaxEntropy = allEntropies.Max();
            result.EntropyStdDev = CalculateStdDev(allEntropies, result.AvgEntropy);
        }

        return result;
    }

    /// <summary>
    /// 快速估算内存区域熵值（单采样点，用于快速扫描模式）
    /// </summary>
    public double QuickEntropyEstimate(IntPtr hProcess, ulong baseAddress, ulong regionSize)
    {
        // 只采样区域中间的 256 字节
        var sampleSize = 256u;
        var offset = regionSize / 2;
        if (offset + sampleSize > regionSize) offset = 0;

        var buffer = new byte[sampleSize];
        if (ReadProcessMemory(hProcess, (IntPtr)(baseAddress + offset), buffer, (int)Math.Min(sampleSize, regionSize), out _))
        {
            return CalculateShannonEntropy(buffer);
        }
        return 0;
    }

    /// <summary>
    /// Shannon 熵计算
    /// </summary>
    public static double CalculateShannonEntropy(byte[] data)
    {
        if (data == null || data.Length == 0) return 0;

        var counts = new int[256];
        foreach (var b in data)
            counts[b]++;

        double entropy = 0;
        double len = data.Length;
        for (int i = 0; i < 256; i++)
        {
            if (counts[i] == 0) continue;
            var p = counts[i] / len;
            entropy -= p * Math.Log2(p);
        }

        return entropy; // 范围 [0, 8]
    }

    private static double CalculateStdDev(List<double> values, double mean)
    {
        if (values.Count <= 1) return 0;
        var sumSq = values.Sum(v => Math.Pow(v - mean, 2));
        return Math.Sqrt(sumSq / (values.Count - 1));
    }

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern bool ReadProcessMemory(IntPtr hProcess, IntPtr lpBaseAddress,
        byte[] lpBuffer, int dwSize, out IntPtr lpNumberOfBytesRead);
}

/// <summary>
/// 熵分析结果
/// </summary>
public class EntropyResult
{
    public double AvgEntropy { get; set; }
    public double MaxEntropy { get; set; }
    public double EntropyStdDev { get; set; }
    public int HighEntropyRegions { get; set; }
    public ulong HighEntropyTotalSize { get; set; }
}