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

namespace MemScanEDR.Services;

/// <summary>
/// MITRE ATT&CK 规则引擎 —— 第三层复核
/// 基于 MITRE ATT&CK 框架中的内存攻击技术，对 AI 判定结果进行交叉验证
/// 降低误报：AI 标记但规则不匹配 → 降级
/// 提高准确率：AI + 规则双重确认 → 确认为恶意
/// </summary>
public class RuleEngine
{
    private readonly List<MitreRule> _rules;

    public RuleEngine()
    {
        _rules = BuildMitreRules();
    }

    /// <summary>
    /// 对 AI 分类结果进行规则复核
    /// </summary>
    public RuleValidationResult Validate(MemoryFeatures features, ClassificationResult aiResult,
        BaselineResult baselineResult)
    {
        var result = new RuleValidationResult();
        var matchedRules = new List<MitreRule>();
        var totalScore = 0.0;

        foreach (var rule in _rules)
        {
            var (matched, confidence) = rule.Evaluate(features);
            if (matched)
            {
                matchedRules.Add(rule);
                totalScore += confidence;
                result.MatchedTechniques.Add(rule.TechniqueId);
            }
        }

        result.MatchedRuleCount = matchedRules.Count;
        result.RuleConfidenceScore = Math.Min(totalScore / Math.Max(_rules.Count, 1), 1.0);

        // ─── 三级交叉验证 ───
        // AI 分数 + 基线异常 + 规则匹配 = 综合判定
        result.AiScore = aiResult.FinalScore;
        result.BaselineAnomalyScore = baselineResult.AnomalyScore;
        result.RuleScore = result.RuleConfidenceScore;

        // 综合分：AI 40% + 基线 30% + 规则 30%
        var combinedScore = aiResult.FinalScore * 0.40 +
                            baselineResult.AnomalyScore * 0.30 +
                            result.RuleConfidenceScore * 0.30;

        result.CombinedScore = Math.Min(combinedScore, 1.0);

        // 最终风险等级
        result.FinalRiskLevel = result.CombinedScore switch
        {
            >= 0.75 => RiskLevel.High,
            >= 0.45 => RiskLevel.Suspicious,
            _ => RiskLevel.Safe
        };

        // ─── 误报修正逻辑 ───
        ApplyFalsePositiveCorrections(result, features);

        // 生成综合解释
        result.Reasoning = BuildReasoning(result, features, matchedRules);

        return result;
    }

    /// <summary>
    /// 误报修正 —— 对AI高分但规则不匹配的情况进行降级
    /// </summary>
    private static void ApplyFalsePositiveCorrections(RuleValidationResult result, MemoryFeatures features)
    {
        // 修正1: JIT 编译器合法 RWX (.NET, Java, V8)
        if (features.RwxPrivatePages > 0 && result.MatchedRuleCount == 0)
        {
            var name = features.ProcessName.ToLowerInvariant();
            var isRuntime = name.Contains("dotnet") || name.Contains("java") ||
                           name.Contains("node") || name.Contains("python") ||
                           name.Contains("chrome") || name.Contains("msedge") ||
                           name.Contains("firefox");

            if (isRuntime && features.IsSystemPath && features.IsSigned)
            {
                // JIT 运行时正常使用 RWX，降级
                result.FinalRiskLevel = RiskLevel.Safe;
                result.CombinedScore = Math.Min(result.CombinedScore, 0.25);
                result.FalsePositiveCorrection = "JIT运行时合法RWX内存（.NET/Java/V8引擎）";
            }
        }

        // 修正2: 安全软件正常行为
        if (features.SuspiciousApiStrings.Count > 0 && result.MatchedRuleCount <= 1)
        {
            var nameLower = features.ProcessName.ToLowerInvariant();
            var isSecurity = nameLower.Contains("defender") || nameLower.Contains("antivirus") ||
                            nameLower.Contains("security") || nameLower.Contains("edr") ||
                            nameLower.Contains("sentinel") || nameLower.Contains("crowd") ||
                            nameLower.Contains("carbon") || nameLower.Contains("mcafee");

            if (isSecurity && features.IsSigned)
            {
                result.FinalRiskLevel = RiskLevel.Safe;
                result.CombinedScore = Math.Min(result.CombinedScore, 0.20);
                result.FalsePositiveCorrection = "安全软件正常行为特征";
                return; // 已经确定是误报，无需继续检查
            }
        }

        // 修正2.5: 脚本引擎/Shell自身字符串误匹配
        if (features.SuspiciousStringHits > 0 && result.MatchedRuleCount <= 1)
        {
            var nameLower = features.ProcessName.ToLowerInvariant();
            var isScriptEngine = nameLower is "powershell.exe" or "pwsh.exe" or "powershell_ise.exe"
                or "cmd.exe" or "wscript.exe" or "cscript.exe" or "conhost.exe"
                or "python.exe" or "python3.exe" or "wsl.exe" or "bash.exe";

            if (isScriptEngine && features.IsSigned && features.IsSystemPath)
            {
                result.FinalRiskLevel = RiskLevel.Safe;
                result.CombinedScore = Math.Min(result.CombinedScore, 0.15);
                result.FalsePositiveCorrection = "脚本引擎自身合法字符串(非恶意注入)";
                return;
            }
        }

        // 修正3: 开发工具调试行为
        if (features.SuspiciousApiStrings.Count > 0 && result.MatchedRuleCount <= 1)
        {
            var exeLower = features.ExePath.ToLowerInvariant();
            var isDev = exeLower.Contains("visual studio") || exeLower.Contains("jetbrains") ||
                       exeLower.Contains("vscode") || exeLower.Contains("debugger") ||
                       exeLower.Contains("ollydbg") || exeLower.Contains("x64dbg") ||
                       exeLower.Contains("windbg") || exeLower.Contains("ida");

            if (isDev)
            {
                result.FinalRiskLevel = RiskLevel.Safe;
                result.CombinedScore = Math.Min(result.CombinedScore, 0.30);
                result.FalsePositiveCorrection = "开发/调试工具正常行为";
            }
        }
    }

    private static string BuildReasoning(RuleValidationResult result, MemoryFeatures f,
        List<MitreRule> matchedRules)
    {
        var parts = new List<string>();

        if (result.FalsePositiveCorrection != null)
        {
            parts.Add($"误报修正: {result.FalsePositiveCorrection}");
        }

        if (matchedRules.Count > 0)
        {
            var ruleDescriptions = matchedRules.Take(3)
                .Select(r => $"{r.TechniqueId} ({r.TechniqueName})");
            parts.Add($"匹配MITRE规则: {string.Join(", ", ruleDescriptions)}");
        }

        if (f.RwxPrivatePages > 0)
            parts.Add($"RWX私有页: {f.RwxPrivatePages}");

        if (f.PeMismatchDisk)
            parts.Add("Process Hollowing 特征");

        if (f.SuspiciousStringHits > 0)
            parts.Add($"敏感字符串命中: {f.SuspiciousStringHits}");

        parts.Add($"综合评分: AI={result.AiScore:F2} 基线={result.BaselineAnomalyScore:F2} 规则={result.RuleScore:F2} → {result.CombinedScore:F2}");

        return string.Join(" | ", parts);
    }

    // ─── MITRE ATT&CK 内存攻击规则库 ───
    private static List<MitreRule> BuildMitreRules() => new()
    {
        // T1055: Process Injection
        new MitreRule
        {
            TechniqueId = "T1055",
            TechniqueName = "进程注入",
            Description = "通过 VirtualAllocEx/WriteProcessMemory/CreateRemoteThread 实现跨进程代码注入",
            EvaluateFunc = f => f.RwxPrivatePages > 0 &&
                                 f.SuspiciousApiStrings.Any(a =>
                                     a.Contains("VirtualAllocEx", StringComparison.OrdinalIgnoreCase) ||
                                     a.Contains("WriteProcessMemory", StringComparison.OrdinalIgnoreCase) ||
                                     a.Contains("CreateRemoteThread", StringComparison.OrdinalIgnoreCase))
        },

        // T1055.012: Process Hollowing
        new MitreRule
        {
            TechniqueId = "T1055.012",
            TechniqueName = "进程掏空 (Process Hollowing)",
            Description = "创建挂起进程后掏空原始代码，替换为恶意载荷",
            EvaluateFunc = f => f.PeMismatchDisk
        },

        // T1055.001: Dynamic-link Library Injection
        new MitreRule
        {
            TechniqueId = "T1055.001",
            TechniqueName = "DLL注入",
            Description = "通过 LoadLibrary/CreateRemoteThread 注入恶意DLL",
            EvaluateFunc = f => f.SuspiciousApiStrings.Any(a =>
                                     a.Contains("LoadLibrary", StringComparison.OrdinalIgnoreCase) ||
                                     a.Contains("LdrLoadDll", StringComparison.OrdinalIgnoreCase)) &&
                                 f.SuspiciousApiStrings.Any(a =>
                                     a.Contains("CreateRemoteThread", StringComparison.OrdinalIgnoreCase))
        },

        // T1620: Reflective Code Loading
        new MitreRule
        {
            TechniqueId = "T1620",
            TechniqueName = "反射代码加载",
            Description = "在内存中直接加载PE/DLL，不通过标准 LoadLibrary",
            EvaluateFunc = f => f.AnonymousExecPages > 0 &&
                                 f.RwxPrivatePages > 0 &&
                                 f.SuspiciousStrings.Any(s =>
                                     s.Contains("ReflectiveLoader", StringComparison.OrdinalIgnoreCase))
        },

        // T1574.002: DLL Side-Loading
        new MitreRule
        {
            TechniqueId = "T1574.002",
            TechniqueName = "DLL侧加载",
            Description = "利用合法程序的DLL搜索顺序加载恶意DLL",
            EvaluateFunc = f => f.TempPathModules > 0 && f.IsSigned && f.IsSystemPath
        },

        // T1071: Application Layer Protocol (C2)
        new MitreRule
        {
            TechniqueId = "T1071",
            TechniqueName = "C2 通信",
            Description = "通过应用层协议与C2服务器通信",
            EvaluateFunc = f => f.ForeignConnections > 0 &&
                                 f.SuspiciousPortConnections > 0 &&
                                 f.SuspiciousApiStrings.Any(a =>
                                     a.Contains("connect", StringComparison.OrdinalIgnoreCase) ||
                                     a.Contains("InternetConnect", StringComparison.OrdinalIgnoreCase))
        },

        // T1095: Non-Application Layer Protocol
        new MitreRule
        {
            TechniqueId = "T1095",
            TechniqueName = "非标准协议通信",
            Description = "使用非标准协议的网络通信",
            EvaluateFunc = f => f.NonStandardPortConnections >= 2 &&
                                 f.ForeignConnections > 0
        },

        // T1027: Obfuscated Files or Information
        new MitreRule
        {
            TechniqueId = "T1027",
            TechniqueName = "代码混淆/加密",
            Description = "内存中存在加密或混淆载荷",
            EvaluateFunc = f => f.MaxEntropy > 7.2 &&
                                 f.HighEntropyRegions >= 2 &&
                                 f.AnonymousExecPages > 0
        },

        // T1003.001: LSASS Memory (Credential Dumping)
        new MitreRule
        {
            TechniqueId = "T1003.001",
            TechniqueName = "LSASS 凭据转储",
            Description = "访问 LSASS 进程内存提取凭据",
            EvaluateFunc = f => f.ProcessName.Equals("lsass.exe", StringComparison.OrdinalIgnoreCase) &&
                                 f.SuspiciousApiStrings.Any(a =>
                                     a.Contains("MiniDump", StringComparison.OrdinalIgnoreCase) ||
                                     a.Contains("SamOpenUser", StringComparison.OrdinalIgnoreCase)) ||
                                 f.SuspiciousStrings.Any(s =>
                                     s.Contains("lsass", StringComparison.OrdinalIgnoreCase))
        },

        // T1036.005: Match Legitimate Name or Location
        new MitreRule
        {
            TechniqueId = "T1036.005",
            TechniqueName = "进程名伪装",
            Description = "使用合法的系统进程名但位于非系统路径",
            EvaluateFunc = f => f.IsSystemProcessName && !f.IsSystemPath
        },

        // T1547.001: Registry Run Keys
        new MitreRule
        {
            TechniqueId = "T1547.001",
            TechniqueName = "注册表持久化",
            Description = "通过注册表run键实现持久化驻留",
            EvaluateFunc = f => f.SuspiciousStrings.Any(s =>
                s.Contains("Run", StringComparison.OrdinalIgnoreCase) &&
                (s.Contains("SOFTWARE\\Microsoft\\Windows\\CurrentVersion", StringComparison.OrdinalIgnoreCase) ||
                 s.Contains("RunOnce", StringComparison.OrdinalIgnoreCase)))
        },

        // T1496: Resource Hijacking (Crypto Mining)
        new MitreRule
        {
            TechniqueId = "T1496",
            TechniqueName = "资源劫持 (挖矿)",
            Description = "劫持系统资源进行加密货币挖矿",
            EvaluateFunc = f => f.SuspiciousStrings.Any(s =>
                s.Contains("stratum", StringComparison.OrdinalIgnoreCase) ||
                s.Contains("cryptonight", StringComparison.OrdinalIgnoreCase) ||
                s.Contains("coinmine", StringComparison.OrdinalIgnoreCase) ||
                s.Contains("mining", StringComparison.OrdinalIgnoreCase))
        },

        // T1486: Data Encrypted for Impact (Ransomware)
        new MitreRule
        {
            TechniqueId = "T1486",
            TechniqueName = "勒索加密",
            Description = "加密文件以勒索赎金",
            EvaluateFunc = f => f.SuspiciousStrings.Any(s =>
                s.Contains("ENCRYPTED", StringComparison.OrdinalIgnoreCase) ||
                s.Contains("ransom", StringComparison.OrdinalIgnoreCase) ||
                s.Contains("decrypt", StringComparison.OrdinalIgnoreCase)) &&
                f.MaxEntropy > 7.0
        },

        // T1562.001: Disable or Modify Tools
        new MitreRule
        {
            TechniqueId = "T1562.001",
            TechniqueName = "禁用安全工具",
            Description = "禁用或修改安全软件",
            EvaluateFunc = f => f.SuspiciousApiStrings.Any(a =>
                a.Contains("AdjustTokenPrivileges", StringComparison.OrdinalIgnoreCase)) &&
                f.IsSystemProcessName
        },

        // T1106: Native API
        new MitreRule
        {
            TechniqueId = "T1106",
            TechniqueName = "Native API 调用",
            Description = "通过NT原生API绕过用户态Hook",
            EvaluateFunc = f => f.SuspiciousApiStrings.Count(a =>
                a.StartsWith("Nt", StringComparison.OrdinalIgnoreCase)) >= 2
        },
    };
}

/// <summary>
/// MITRE ATT&CK 单条规则定义
/// </summary>
public class MitreRule
{
    public string TechniqueId { get; init; } = "";
    public string TechniqueName { get; init; } = "";
    public string Description { get; init; } = "";
    public Func<MemoryFeatures, bool> EvaluateFunc { get; init; } = _ => false;

    /// <summary>
    /// 评估规则，返回 (是否匹配, 置信度)
    /// </summary>
    public (bool matched, double confidence) Evaluate(MemoryFeatures features)
    {
        try
        {
            var match = EvaluateFunc(features);
            return (match, match ? 1.0 : 0.0);
        }
        catch
        {
            return (false, 0);
        }
    }
}

/// <summary>
/// 规则引擎验证结果
/// </summary>
public class RuleValidationResult
{
    public double AiScore { get; set; }
    public double BaselineAnomalyScore { get; set; }
    public double RuleScore { get; set; }
    public double CombinedScore { get; set; }
    public RiskLevel FinalRiskLevel { get; set; } = RiskLevel.Safe;
    public int MatchedRuleCount { get; set; }
    public List<string> MatchedTechniques { get; set; } = new();
    public double RuleConfidenceScore { get; set; }
    public string Reasoning { get; set; } = "";
    public string? FalsePositiveCorrection { get; set; }
}