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

namespace MemScanEDR.Services;

/// <summary>
/// 监督学习分类引擎 —— 第二层检测
/// 使用轻量级加权特征分类器，模拟训练好的监督模型的推理行为
/// 不依赖深度学习框架，纯 C# 实现，适合终端 EDR 部署
/// 
/// 分类逻辑基于：
/// 1. 特征加权评分（模拟训练好的线性模型权重）
/// 2. 特征组合规则（模拟决策树的关键分支）
/// 3. 已知恶意模式匹配
/// </summary>
public class ClassificationEngine
{
    // ─── 特征权重（从训练数据中学习到的权重，模拟线性分类器）───
    // 正权重 = 该特征越高越可疑，负权重 = 正常特征
    private static readonly double[] FeatureWeights = new double[]
    {
        // 进程元数据
        -2.0, -1.5,  3.0,  0.1,  0.1,  2.5, -2.0,  0.1,  0.0,
        // 内存页权限（RWX 相关权重极高）
         0.1, -0.5, -0.3, -1.0,  4.0,  3.5,  2.5,
        // 内存类型
        -0.5, -0.3, -1.0,
        // 代码特征
         1.5,  4.0,  2.0,  2.5,  1.5,  2.0,
        // 熵特征
         1.0,  2.5,  2.0,  1.0,
        // 网络特征
         0.5,  0.3,  1.0,  2.5,  1.5,
        // DLL加载
         0.2,  2.0,  2.5
    };

    // 分类阈值
    private const double HighRiskThreshold = 0.70;
    private const double SuspiciousThreshold = 0.40;

    /// <summary>
    /// 对特征向量进行分类
    /// </summary>
    public ClassificationResult Classify(MemoryFeatures features)
    {
        var result = new ClassificationResult();
        var vector = features.ToFeatureVector();

        // ─── 1. 线性加权评分 ───
        double weightedScore = 0;
        double maxPossibleScore = 0;
        var topContributors = new List<(int index, double contribution)>();

        for (int i = 0; i < Math.Min(vector.Length, FeatureWeights.Length); i++)
        {
            var contribution = vector[i] * FeatureWeights[i];
            weightedScore += contribution;

            if (FeatureWeights[i] > 0)
                maxPossibleScore += Math.Abs(FeatureWeights[i]);

            if (Math.Abs(contribution) > 0.5)
                topContributors.Add((i, contribution));
        }

        // 归一化到 [0, 1]
        var normalizedScore = maxPossibleScore > 0
            ? Math.Max(0, weightedScore) / (maxPossibleScore * 0.5)
            : 0;
        normalizedScore = Math.Min(normalizedScore, 1.0);

        // ─── 2. 决策树规则增强 ───
        var ruleBoost = ApplyDecisionRules(features);

        // 综合得分
        var finalScore = normalizedScore * 0.7 + ruleBoost * 0.3;
        finalScore = Math.Min(finalScore, 1.0);

        result.RawScore = normalizedScore;
        result.RuleBoost = ruleBoost;
        result.FinalScore = finalScore;

        // ─── 3. 分类判定 ───
        result.Classification = finalScore switch
        {
            >= HighRiskThreshold => MalwareClass.Malicious,
            >= SuspiciousThreshold => MalwareClass.Suspicious,
            _ => MalwareClass.Benign
        };

        // ─── 4. 恶意家族推测 ───
        result.SuspectedFamily = InferMalwareFamily(features);

        // ─── 5. 关键特征贡献 ───
        result.TopContributors = topContributors
            .OrderByDescending(c => Math.Abs(c.contribution))
            .Take(5)
            .Select(c => (GetFeatureName(c.index), c.contribution))
            .ToList();

        return result;
    }

    /// <summary>
    /// 决策树规则 —— 模拟关键分支逻辑
    /// </summary>
    private static double ApplyDecisionRules(MemoryFeatures f)
    {
        double boost = 0;

        // 规则1: 系统进程 + RWX私有页 = 极高风险 (代码注入核心特征)
        if (f.IsSystemProcessName && f.RwxPrivatePages > 0)
            boost += 0.40;

        // 规则2: 未签名 + 临时路径 + RWX = 高风险
        if (!f.IsSigned && f.IsTempPath && f.RwxPrivatePages > 0)
            boost += 0.30;

        // 规则3: Process Hollowing = 立即高风险
        if (f.PeMismatchDisk)
            boost += 0.50;

        // 规则4: 高熵 + 无磁盘映射可执行页 = 加密载荷
        if (f.MaxEntropy > 7.0 && f.AnonymousExecPages > 0 && f.ImagePages == 0)
            boost += 0.35;

        // 规则5: 签名匹配 + 网络C2端口 = 远控木马
        if (f.HighSeverityMatches > 0 && f.SuspiciousPortConnections > 0)
            boost += 0.35;

        // 规则6: 系统进程名 + 非系统路径 = 进程伪装
        if (f.IsSystemProcessName && !f.IsSystemPath)
            boost += 0.25;

        // 规则7: 可疑API多 + 网络连接 = 可能C2
        if (f.SuspiciousApiStrings.Count >= 3 && f.ConnectionCount > 0)
            boost += 0.20;

        // 规则8: 境外连接 + 非标准端口 + 高熵 = 数据外泄
        if (f.ForeignConnections > 0 && f.NonStandardPortConnections > 0 && f.MaxEntropy > 7.0)
            boost += 0.25;

        // 规则9: 无磁盘模块 + 临时路径模块
        if (f.UnmappedModules > 0 && f.TempPathModules > 0)
            boost += 0.20;

        return Math.Min(boost, 0.60);
    }

    /// <summary>
    /// 推测恶意软件家族
    /// </summary>
    private static string InferMalwareFamily(MemoryFeatures features)
    {
        var indicators = new List<string>();

        // CobaltStrike
        if (features.SuspiciousStrings.Any(s => s.Contains("beacon", StringComparison.OrdinalIgnoreCase)))
            indicators.Add("CobaltStrike Beacon");

        // Meterpreter
        if (features.SuspiciousStrings.Any(s => s.Contains("meterpreter", StringComparison.OrdinalIgnoreCase)))
            indicators.Add("Metasploit Meterpreter");

        // Mimikatz
        if (features.SuspiciousStrings.Any(s => s.Contains("mimikatz", StringComparison.OrdinalIgnoreCase)) ||
            features.SuspiciousApiStrings.Any(a => a.Contains("MiniDump", StringComparison.OrdinalIgnoreCase)))
            indicators.Add("凭证窃取工具 (Mimikatz)");

        // Process Hollowing
        if (features.PeMismatchDisk)
            indicators.Add("Process Hollowing");

        // Ransomware
        if (features.SuspiciousStrings.Any(s => s.Contains("ENCRYPTED", StringComparison.OrdinalIgnoreCase) ||
                                                s.Contains("ransom", StringComparison.OrdinalIgnoreCase)))
            indicators.Add("勒索软件");

        // Crypto Miner
        if (features.SuspiciousStrings.Any(s => s.Contains("stratum", StringComparison.OrdinalIgnoreCase) ||
                                                s.Contains("cryptonight", StringComparison.OrdinalIgnoreCase) ||
                                                s.Contains("coinmine", StringComparison.OrdinalIgnoreCase)))
            indicators.Add("挖矿木马");

        // RAT
        if (features.SuspiciousApiStrings.Count >= 3 && features.ForeignConnections > 0 &&
            features.SuspiciousPortConnections > 0)
            indicators.Add("远控木马 (RAT)");

        // Reflective DLL Injection
        if (features.AnonymousExecPages > 0 && features.RwxPrivatePages > 0 &&
            features.SuspiciousApiStrings.Any(a => a.Contains("Reflective", StringComparison.OrdinalIgnoreCase)))
            indicators.Add("反射DLL注入木马");

        // Bootkit/Rootkit
        if (features.IsSystemProcessName && !features.IsSystemPath && features.RwxPrivatePages > 0)
            indicators.Add("Rootkit/伪装系统进程");

        return indicators.Count > 0 ? string.Join(" | ", indicators) : "未知恶意家族";
    }

    private static string GetFeatureName(int index) => index switch
    {
        0 => "未签名", 1 => "非系统路径", 2 => "临时路径",
        3 => "内存占用", 4 => "线程数", 5 => "系统进程名", 6 => "父进程异常",
        12 => "RWX页比例", 13 => "匿名可执行页", 14 => "RWX私有页",
        18 => "内存PE头", 19 => "PE与磁盘不一致",
        20 => "可疑API", 21 => "敏感字符串", 22 => "签名匹配", 23 => "高危签名",
        24 => "高熵代码段", 25 => "最大熵值", 26 => "高熵区域数",
        29 => "非标准端口", 30 => "可疑端口", 31 => "境外连接",
        35 => "无磁盘模块", 36 => "临时目录模块",
        _ => $"特征{index}"
    };
}

/// <summary>
/// 分类引擎结果
/// </summary>
public class ClassificationResult
{
    public MalwareClass Classification { get; set; } = MalwareClass.Benign;
    public double RawScore { get; set; }
    public double RuleBoost { get; set; }
    public double FinalScore { get; set; }
    public string SuspectedFamily { get; set; } = "";
    public List<(string featureName, double contribution)> TopContributors { get; set; } = new();
}

public enum MalwareClass
{
    Benign,
    Suspicious,
    Malicious
}