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

namespace MemScanEDR.Services;

/// <summary>
/// 特征提取器 —— 核心模块
/// 从进程内存中提取全部多维特征，作为 AI 模型的输入向量
/// 包含：元数据、页权限分布、代码特征、熵值、网络行为、进程树关系
/// </summary>
public class FeatureExtractor
{
    private readonly EntropyAnalyzer _entropyAnalyzer = new();
    private readonly PEHeaderAnalyzer _peAnalyzer = new();

    private const uint PROCESS_VM_READ = 0x0010;
    private const uint PROCESS_QUERY_INFORMATION = 0x0400;

    /// <summary>
    /// 从进程信息中提取完整的 MemoryFeatures 特征向量
    /// </summary>
    public MemoryFeatures Extract(ProcessInfo proc, List<MemoryRegion> regions,
        List<SignatureMatch> matches, List<Signature> signatures)
    {
        var features = new MemoryFeatures();

        // ─── 一、进程元数据 ───
        features.ProcessName = proc.Name;
        features.ExePath = proc.ExePath;
        features.ParentPid = proc.Ppid;
        features.ParentName = proc.ParentName;
        features.CmdLine = proc.CmdLine;
        features.StartTime = proc.StartTime;
        features.WorkingSetMB = proc.WorkingSetMB;
        features.ThreadCount = proc.ThreadCount;
        features.IsSigned = proc.IsSigned;
        features.Signer = proc.Signer;
        features.IsSystemPath = proc.IsSystemPath;
        features.IsTempPath = proc.IsTempPath;

        // ─── 二、内存页权限分布 ───
        ExtractPagePermissions(features, regions);

        // ─── 三、内存代码特征 ───
        ExtractCodeFeatures(features, proc, regions, matches);

        // ─── 四、熵值特征 ───
        ExtractEntropyFeatures(features, proc.Pid, regions);

        // ─── 五、网络行为特征 ───
        ExtractNetworkFeatures(features, proc.Connections);

        // ─── 六、进程树关系 ───
        features.IsSystemProcessName = IsSystemProcessName(proc.Name);
        features.HasLegitimateParent = HasLegitimateParent(proc.ParentName, proc.Name);

        // ─── 七、DLL 加载特征 ───
        features.LoadedModuleCount = proc.LoadedModules.Count;
        features.UnmappedModules = CountUnmappedModules(regions);
        features.TempPathModules = CountTempModules(proc.LoadedModules);

        return features;
    }

    /// <summary>
    /// 快速特征提取 —— 仅提取元数据和权限分布 (不读内存，速度最快)
    /// 用于快速预过滤阶段
    /// </summary>
    public MemoryFeatures QuickExtract(ProcessInfo proc, List<MemoryRegion> regions)
    {
        var features = new MemoryFeatures
        {
            ProcessName = proc.Name,
            ExePath = proc.ExePath,
            ParentPid = proc.Ppid,
            ParentName = proc.ParentName,
            IsSigned = proc.IsSigned,
            IsSystemPath = proc.IsSystemPath,
            IsTempPath = proc.IsTempPath,
            WorkingSetMB = proc.WorkingSetMB,
            ThreadCount = proc.ThreadCount,
            IsSystemProcessName = IsSystemProcessName(proc.Name),
            HasLegitimateParent = HasLegitimateParent(proc.ParentName, proc.Name)
        };

        ExtractPagePermissions(features, regions);
        ExtractNetworkFeatures(features, proc.Connections);

        return features;
    }

    // ─── 内存页权限提取 ───
    private static void ExtractPagePermissions(MemoryFeatures f, List<MemoryRegion> regions)
    {
        f.TotalRegions = regions.Count;

        foreach (var r in regions)
        {
            switch (r.Protect)
            {
                case "READONLY": f.ReadOnlyPages++; break;
                case "READWRITE": f.ReadWritePages++; break;
                case "EXECUTE": case "EXECUTE_READ": f.ExecuteReadPages++; break;
                case "EXECUTE_READWRITE": f.ExecuteReadWritePages++; break;
            }

            switch (r.Type)
            {
                case "PRIVATE": f.PrivatePages++; break;
                case "MAPPED": f.MappedPages++; break;
                case "IMAGE": f.ImagePages++; break;
            }

            // 无文件匿名可执行内存
            if (r.Type != "IMAGE" && (r.Protect.Contains("EXECUTE")))
            {
                f.AnonymousExecPages++;
            }

            // RWX 私有页 —— 核心恶意指标
            if (r.Protect.Contains("EXECUTE_READWRITE") && r.Type == "PRIVATE")
            {
                f.RwxPrivatePages++;
                f.RwxTotalSize += r.Size;
            }
        }
    }

    // ─── 代码特征提取 ───
    private void ExtractCodeFeatures(MemoryFeatures f, ProcessInfo proc,
        List<MemoryRegion> regions, List<SignatureMatch> matches)
    {
        // 脚本引擎/运行时自身是合法进程，跳过内存字符串扫描避免误报
        if (IsScriptingEngine(proc))
        {
            f.SignatureMatchCount = matches.Count;
            f.HighSeverityMatches = matches.Count(m => m.Severity == "high");
            return;
        }

        // 打开进程读取内存
        var hProcess = OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, false, proc.Pid);
        if (hProcess == IntPtr.Zero) return;

        try
        {
            // PE 头检测
            foreach (var region in regions.Where(r => r.Type == "IMAGE" && r.State == "COMMIT"))
            {
                if (_peAnalyzer.HasInMemoryPE(hProcess, region.BaseAddress))
                {
                    f.HasInMemoryPE = true;

                    // 检查内存PE与磁盘是否一致
                    if (!string.IsNullOrEmpty(proc.ExePath))
                    {
                        var integrity = _peAnalyzer.CheckPeIntegrity(proc.ExePath, hProcess, region.BaseAddress);
                        f.PeMismatchDisk = !integrity.IsValid;
                    }
                    break;
                }
            }

            // 敏感字符串提取 (只在可疑区域提取以提升速度)
            var suspiciousRegions = regions.Where(r =>
                r.Type == "PRIVATE" && r.Protect.Contains("EXECUTE") && r.Size < 10 * 1024 * 1024
            ).Take(5);

            foreach (var region in suspiciousRegions)
            {
                var strings = _peAnalyzer.ExtractSuspiciousStrings(hProcess, region.BaseAddress, region.Size);
                f.SuspiciousStrings.AddRange(strings.Take(5));
                f.SuspiciousStringHits += strings.Count;

                var apis = _peAnalyzer.ExtractSuspiciousApis(hProcess, region.BaseAddress, Math.Min(region.Size, 32768));
                f.SuspiciousApiStrings.AddRange(apis);
            }

            // 去重
            f.SuspiciousStrings = f.SuspiciousStrings.Distinct().ToList();
            f.SuspiciousApiStrings = f.SuspiciousApiStrings.Distinct().ToList();
        }
        finally
        {
            CloseHandle(hProcess);
        }

        // 签名匹配
        f.SignatureMatchCount = matches.Count;
        f.HighSeverityMatches = matches.Count(m => m.Severity == "high");
    }

    // ─── 熵值特征提取 ───
    private void ExtractEntropyFeatures(MemoryFeatures f, int pid, List<MemoryRegion> regions)
    {
        var hProcess = OpenProcess(PROCESS_VM_READ, false, pid);
        if (hProcess == IntPtr.Zero) return;

        try
        {
            var result = _entropyAnalyzer.AnalyzeRegions(hProcess, regions);
            f.AvgCodeEntropy = result.AvgEntropy;
            f.MaxEntropy = result.MaxEntropy;
            f.HighEntropyRegions = result.HighEntropyRegions;
            f.HighEntropyTotalSize = result.HighEntropyTotalSize;
            f.EntropyStdDev = result.EntropyStdDev;
        }
        finally
        {
            CloseHandle(hProcess);
        }
    }

    // ─── 网络特征提取 ───
    private static void ExtractNetworkFeatures(MemoryFeatures f, List<NetworkConnection> connections)
    {
        f.ConnectionCount = connections.Count;

        foreach (var conn in connections)
        {
            if (conn.State == "Listen")
                f.ListeningPorts++;

            // 非标准端口
            if (conn.RemotePort > 0 && conn.RemotePort != 80 && conn.RemotePort != 443
                && conn.RemotePort != 8080 && conn.RemotePort != 8443)
            {
                f.NonStandardPortConnections++;
            }

            // 已知C2端口
            if (IsSuspiciousPort(conn.RemotePort))
                f.SuspiciousPortConnections++;

            // 境外连接检测 (简化判断)
            if (IsForeignAddress(conn.RemoteAddr))
                f.ForeignConnections++;
        }
    }

    private static bool IsSuspiciousPort(int port) => port switch
    {
        4444 or 5555 or 6666 or 7777 or 1337 or 31337 or 9999 or 12345 or 27015 => true,
        _ => false
    };

    private static bool IsForeignAddress(string addr)
    {
        if (string.IsNullOrEmpty(addr)) return false;
        // 私有地址范围
        var ip = addr.Split(':')[0];
        return !(ip.StartsWith("10.") || ip.StartsWith("172.16.") ||
                 ip.StartsWith("192.168.") || ip == "127.0.0.1" || ip == "0.0.0.0");
    }

    // ─── 进程树辅助方法 ───
    public static bool IsSystemProcessName(string name) =>
        new HashSet<string>(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"
        }.Contains(name);

    public static bool HasLegitimateParent(string parentName, string childName)
    {
        // 正常父子关系
        if (string.IsNullOrEmpty(parentName)) return false;

        var childLower = childName.ToLowerInvariant();
        var parentLower = parentName.ToLowerInvariant();

        // 系统进程的正常父进程
        if (parentLower is "services.exe" or "wininit.exe" or "smss.exe")
            return true;

        // Explorer 启动用户程序
        if (parentLower == "explorer.exe")
            return true;

        // 浏览器多进程
        if ((parentLower.Contains("chrome") || parentLower.Contains("msedge") || parentLower.Contains("firefox"))
            && (childLower.Contains("chrome") || childLower.Contains("msedge") || childLower.Contains("firefox")))
            return true;

        // svchost 只由 services.exe 启动
        if (childLower == "svchost.exe" && parentLower == "services.exe")
            return true;

        return false;
    }

    /// <summary>
    /// 判断是否为脚本引擎/Shell —— 这类进程自身包含大量"可疑"字符串
    /// 但属于正常功能（如 PowerShell 的 Invoke-Expression, IEX 等是合法的内置cmdlet）
    /// </summary>
    public static bool IsScriptingEngine(ProcessInfo proc)
    {
        if (!proc.IsSigned || !proc.IsSystemPath) return false;

        var nameLower = proc.Name.ToLowerInvariant();
        return nameLower switch
        {
            "powershell.exe" or "pwsh.exe" or "powershell_ise.exe" => true,
            "cmd.exe" => true,
            "wscript.exe" or "cscript.exe" => true,
            "conhost.exe" => true,
            "python.exe" or "python3.exe" or "pythonw.exe" => true,
            "wsl.exe" or "bash.exe" => true,
            _ => false
        };
    }

    private static int CountUnmappedModules(List<MemoryRegion> regions) =>
        regions.Count(r => r.Type == "PRIVATE" && r.Protect.Contains("EXECUTE") && r.Size > 4096);

    private static int CountTempModules(List<string> modules) =>
        modules.Count(m => m.Contains("\\Temp\\", StringComparison.OrdinalIgnoreCase) ||
                           m.Contains("\\tmp\\", StringComparison.OrdinalIgnoreCase) ||
                           m.Contains("\\AppData\\Local\\Temp\\", StringComparison.OrdinalIgnoreCase));

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

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