package com.aluer.defense;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicLong;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.regex.Pattern;
import java.util.regex.Matcher;

@Component
public class AdvancedMalwareDetectionService {
    private static final Logger logger = LoggerFactory.getLogger(AdvancedMalwareDetectionService.class);

    private final Map<String, MalwareSignature> signatures = new ConcurrentHashMap<>();
    private final Map<String, BehavioralProfile> behavioralProfiles = new ConcurrentHashMap<>();
    private final Queue<MalwareAlert> alertQueue = new ConcurrentLinkedQueue<>();
    private final Map<String, List<HeuristicResult>> scanHistory = new ConcurrentHashMap<>();
    private final AtomicLong totalScans = new AtomicLong(0);
    private final AtomicLong threatsDetected = new AtomicLong(0);

    private static final int MAX_SCAN_HISTORY = 1000;
    private static final DateTimeFormatter DATE_FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");

    public AdvancedMalwareDetectionService() {
        initializeSignatures();
        logger.info("Advanced Malware Detection Service initialized");
    }

    private void initializeSignatures() {
        addSignature("WIN32_TROJAN_GENERIC", "Trojan.Generic", "HIGH", "eicar_test_file", "Known trojan signature");
        addSignature("WIN32_WORM_CONFICKER", "Worm.Conficker", "CRITICAL", "conficker_.*", "Conficker worm");
        addSignature("WIN32_RANSOMWARE_CRYPTO", "Ransomware.CryptoLocker", "CRITICAL", "crypto_locker", "CryptoLocker ransomware");
        addSignature("WIN32_BACKDOOR_TITAN", "Backdoor.Titan", "HIGH", "titan_.*", "Titan backdoor");
        addSignature("WIN64_ROOTKIT_FLYPRO", "Rootkit.FlyPro", "HIGH", "flypro_.*", "FlyPro rootkit");
        addSignature("WIN32_ADWARE_YOUNEE", "Adware.YooNear", "LOW", "yoonee_.*", "Adware component");
        addSignature("WIN32_KEYLOGGER_SPYWARE", "Keylogger.SpyWare", "HIGH", "keylog.*", "Keylogger detection");
        addSignature("WIN32_BOTNET_ZOMBIE", "Botnet.Zombie", "HIGH", "zombie_.*", "Botnet zombie");
        addSignature("WIN32_MINER_COIN", "Miner.CoinHive", "MEDIUM", "coinhive.*", "Cryptocurrency miner");
        addSignature("WIN32_EXPLOIT_KIT", "ExploitKit.EK", "HIGH", "exploit_.*", "Exploit kit");

        addSignature("LINUX_MALWARE_MIRAI", "Linux.Mirai", "CRITICAL", "mirai_.*", "Mirai botnet");
        addSignature("LINUX_ROOTKIT_DIAMONDER", "Linux.Rootkit.Diamond", "HIGH", "diamond_.*", "Diamond rootkit");
        addSignature("LINUX_MINER_LINUX", "Linux.Miner.Xmrig", "MEDIUM", "xmrig.*", "XMRig miner");
        addSignature("LINUX_BACKDOOR_BEAST", "Linux.Backdoor.Beast", "HIGH", "beast_.*", "Beast backdoor");

        logger.info("Initialized {} malware signatures", signatures.size());
    }

    public void addSignature(String id, String name, String severity, String pattern, String description) {
        MalwareSignature sig = new MalwareSignature(id, name, severity, pattern, description);
        signatures.put(id, sig);
    }

    public ScanResult scanFile(String fileName, byte[] fileContent, String fileHash) {
        ScanResult result = new ScanResult();
        result.setFileName(fileName);
        result.setFileHash(fileHash);
        result.setTimestamp(LocalDateTime.now());

        totalScans.incrementAndGet();

        MalwareSignature matchedSig = matchSignature(fileHash, fileName);
        if (matchedSig != null) {
            result.setThreatDetected(true);
            result.setThreatName(matchedSig.getName());
            result.setThreatSeverity(matchedSig.getSeverity());
            result.setThreatDescription(matchedSig.getDescription());
            threatsDetected.incrementAndGet();

            triggerAlert(fileName, matchedSig);
            logger.warn("MALWARE DETECTED: {} in file {}", matchedSig.getName(), fileName);
            return result;
        }

        List<HeuristicResult> heuristics = runHeuristicAnalysis(fileContent, fileName);
        for (HeuristicResult hr : heuristics) {
            if (hr.isSuspicious()) {
                result.addHeuristicResult(hr);
                result.setThreatDetected(true);
                result.setThreatName("Heuristic Detection");
                result.setThreatSeverity("MEDIUM");
                threatsDetected.incrementAndGet();
            }
        }

        List<HeuristicResult> history = scanHistory.computeIfAbsent(fileHash, k -> new ArrayList<>());
        history.addAll(heuristics);
        if (history.size() > MAX_SCAN_HISTORY) {
            history.subList(0, history.size() - MAX_SCAN_HISTORY).clear();
        }

        return result;
    }

    private MalwareSignature matchSignature(String fileHash, String fileName) {
        for (MalwareSignature sig : signatures.values()) {
            try {
                Pattern pattern = Pattern.compile(sig.getPattern(), Pattern.CASE_INSENSITIVE);
                if (pattern.matcher(fileHash).matches() || pattern.matcher(fileName).matches()) {
                    return sig;
                }
            } catch (Exception e) {
            }
        }
        return null;
    }

    private List<HeuristicResult> runHeuristicAnalysis(byte[] content, String fileName) {
        List<HeuristicResult> results = new ArrayList<>();

        if (content == null || content.length == 0) {
            return results;
        }

        HeuristicResult entropyResult = analyzeEntropy(content);
        if (entropyResult != null) {
            results.add(entropyResult);
        }

        HeuristicResult packingResult = analyzePacking(content);
        if (packingResult != null) {
            results.add(packingResult);
        }

        HeuristicResult suspiciousStringsResult = analyzeSuspiciousStrings(content);
        if (suspiciousStringsResult != null) {
            results.add(suspiciousStringsResult);
        }

        HeuristicResult suspiciousPatternsResult = analyzeSuspiciousPatterns(content);
        if (suspiciousPatternsResult != null) {
            results.add(suspiciousPatternsResult);
        }

        return results;
    }

    private HeuristicResult analyzeEntropy(byte[] content) {
        double entropy = calculateEntropy(content);

        HeuristicResult result = new HeuristicResult();
        result.setType("ENTROPY");
        result.setScore((int)(entropy * 100));

        if (entropy > 7.5) {
            result.setSuspicious(true);
            result.setReason("Very high entropy: " + String.format("%.2f", entropy) + " (possible encryption or packing)");
        } else if (entropy > 6.5) {
            result.setSuspicious(true);
            result.setReason("High entropy: " + String.format("%.2f", entropy) + " (possible obfuscation)");
        }

        return result;
    }

    private double calculateEntropy(byte[] data) {
        if (data == null || data.length == 0) return 0;

        int[] frequency = new int[256];
        for (byte b : data) {
            frequency[b & 0xFF]++;
        }

        double entropy = 0;
        int length = data.length;
        for (int count : frequency) {
            if (count > 0) {
                double probability = (double) count / length;
                entropy -= probability * (Math.log(probability) / Math.log(2));
            }
        }

        return entropy;
    }

    private HeuristicResult analyzePacking(byte[] content) {
        HeuristicResult result = new HeuristicResult();
        result.setType("PACKING");

        int nullBytes = 0;
        int nonAscii = 0;

        for (int i = 0; i < Math.min(content.length, 10000); i++) {
            if (content[i] == 0) nullBytes++;
            if ((content[i] & 0xFF) > 127) nonAscii++;
        }

        double nullRatio = (double) nullBytes / content.length;
        double nonAsciiRatio = (double) nonAscii / content.length;

        if (nullRatio > 0.3 && nullRatio < 0.7) {
            result.setSuspicious(true);
            result.setScore(70);
            result.setReason("Suspicious null byte ratio: " + String.format("%.2f", nullRatio * 100) + "%");
        } else if (nonAsciiRatio > 0.3) {
            result.setSuspicious(true);
            result.setScore(60);
            result.setReason("High ratio of non-ASCII bytes: " + String.format("%.2f", nonAsciiRatio * 100) + "%");
        }

        return result;
    }

    private HeuristicResult analyzeSuspiciousStrings(byte[] content) {
        String[] suspiciousPatterns = {
            "http://", "https://", "ftp://",
            "cmd.exe", "powershell", "bash", "/bin/sh",
            "reg add", "reg delete", "regedit",
            "CreateRemoteThread", "VirtualAlloc", "WriteProcessMemory",
            "LoadLibrary", "GetProcAddress", "WinExec",
            "http.request", "http.response",
            "eval(", "exec(", "system(",
            "base64_decode", "gzinflate", "str_rot13"
        };

        HeuristicResult result = new HeuristicResult();
        result.setType("SUSPICIOUS_STRINGS");
        int matchCount = 0;

        String contentStr = new String(content);
        for (String pattern : suspiciousPatterns) {
            if (contentStr.toLowerCase().contains(pattern.toLowerCase())) {
                matchCount++;
            }
        }

        if (matchCount > 5) {
            result.setSuspicious(true);
            result.setScore(Math.min(100, matchCount * 10));
            result.setReason("Found " + matchCount + " suspicious string patterns");
        }

        return result;
    }

    private HeuristicResult analyzeSuspiciousPatterns(byte[] content) {
        HeuristicResult result = new HeuristicResult();
        result.setType("SUSPICIOUS_PATTERNS");
        int score = 0;

        if (content.length > 1024 * 1024) {
            score += 30;
        }

        int consecutiveRepeating = 0;
        for (int i = 1; i < Math.min(content.length, 10000); i++) {
            if (content[i] == content[i-1]) {
                consecutiveRepeating++;
            } else {
                consecutiveRepeating = 0;
            }
            if (consecutiveRepeating > 100) {
                score += 20;
                break;
            }
        }

        if (score > 0) {
            result.setSuspicious(true);
            result.setScore(score);
            result.setReason("Suspicious binary patterns detected");
        }

        return result;
    }

    public void updateBehavioralProfile(String identifier, String behavior, int severity) {
        BehavioralProfile profile = behavioralProfiles.computeIfAbsent(identifier, k -> new BehavioralProfile(identifier));
        profile.addBehavior(behavior, severity);
    }

    private void triggerAlert(String fileName, MalwareSignature signature) {
        MalwareAlert alert = new MalwareAlert(
            fileName,
            signature.getName(),
            signature.getSeverity(),
            signature.getDescription(),
            LocalDateTime.now()
        );
        alertQueue.offer(alert);

        if (alertQueue.size() > 1000) {
            alertQueue.poll();
        }
    }

    public Map<String, Object> getStats() {
        Map<String, Object> stats = new HashMap<>();
        stats.put("totalScans", totalScans.get());
        stats.put("threatsDetected", threatsDetected.get());
        stats.put("signaturesLoaded", signatures.size());
        stats.put("activeProfiles", behavioralProfiles.size());
        stats.put("alertQueueSize", alertQueue.size());

        long critical = signatures.values().stream().filter(s -> "CRITICAL".equals(s.getSeverity())).count();
        long high = signatures.values().stream().filter(s -> "HIGH".equals(s.getSeverity())).count();
        stats.put("criticalSignatures", critical);
        stats.put("highSignatures", high);

        return stats;
    }

    public List<MalwareAlert> getRecentAlerts(int limit) {
        List<MalwareAlert> alerts = new ArrayList<>();
        int count = 0;
        for (MalwareAlert alert : alertQueue) {
            if (count++ >= limit) break;
            alerts.add(alert);
        }
        return alerts;
    }

    public Map<String, BehavioralProfile> getBehavioralProfiles() {
        return new HashMap<>(behavioralProfiles);
    }

    public static class MalwareSignature {
        private final String id;
        private final String name;
        private final String severity;
        private final String pattern;
        private final String description;

        public MalwareSignature(String id, String name, String severity, String pattern, String description) {
            this.id = id;
            this.name = name;
            this.severity = severity;
            this.pattern = pattern;
            this.description = description;
        }

        public String getId() { return id; }
        public String getName() { return name; }
        public String getSeverity() { return severity; }
        public String getPattern() { return pattern; }
        public String getDescription() { return description; }
    }

    public static class ScanResult {
        private String fileName;
        private String fileHash;
        private boolean threatDetected;
        private String threatName;
        private String threatSeverity;
        private String threatDescription;
        private List<HeuristicResult> heuristicResults;
        private LocalDateTime timestamp;

        public ScanResult() {
            this.heuristicResults = new ArrayList<>();
        }

        public String getFileName() { return fileName; }
        public void setFileName(String fileName) { this.fileName = fileName; }
        public String getFileHash() { return fileHash; }
        public void setFileHash(String fileHash) { this.fileHash = fileHash; }
        public boolean isThreatDetected() { return threatDetected; }
        public void setThreatDetected(boolean threatDetected) { this.threatDetected = threatDetected; }
        public String getThreatName() { return threatName; }
        public void setThreatName(String threatName) { this.threatName = threatName; }
        public String getThreatSeverity() { return threatSeverity; }
        public void setThreatSeverity(String threatSeverity) { this.threatSeverity = threatSeverity; }
        public String getThreatDescription() { return threatDescription; }
        public void setThreatDescription(String threatDescription) { this.threatDescription = threatDescription; }
        public List<HeuristicResult> getHeuristicResults() { return heuristicResults; }
        public void addHeuristicResult(HeuristicResult result) { this.heuristicResults.add(result); }
        public LocalDateTime getTimestamp() { return timestamp; }
        public void setTimestamp(LocalDateTime timestamp) { this.timestamp = timestamp; }
    }

    public static class HeuristicResult {
        private String type;
        private int score;
        private boolean suspicious;
        private String reason;

        public String getType() { return type; }
        public void setType(String type) { this.type = type; }
        public int getScore() { return score; }
        public void setScore(int score) { this.score = score; }
        public boolean isSuspicious() { return suspicious; }
        public void setSuspicious(boolean suspicious) { this.suspicious = suspicious; }
        public String getReason() { return reason; }
        public void setReason(String reason) { this.reason = reason; }
    }

    public static class BehavioralProfile {
        private final String identifier;
        private final Map<String, Integer> behaviors;
        private volatile int threatScore;
        private volatile long lastUpdate;

        public BehavioralProfile(String identifier) {
            this.identifier = identifier;
            this.behaviors = new ConcurrentHashMap<>();
            this.lastUpdate = System.currentTimeMillis();
        }

        public void addBehavior(String behavior, int severity) {
            behaviors.merge(behavior, severity, Integer::sum);
            threatScore += severity;
            lastUpdate = System.currentTimeMillis();
        }

        public String getIdentifier() { return identifier; }
        public Map<String, Integer> getBehaviors() { return behaviors; }
        public int getThreatScore() { return threatScore; }
        public long getLastUpdate() { return lastUpdate; }
    }

    public static class MalwareAlert {
        private final String fileName;
        private final String threatName;
        private final String severity;
        private final String description;
        private final LocalDateTime timestamp;

        public MalwareAlert(String fileName, String threatName, String severity, String description, LocalDateTime timestamp) {
            this.fileName = fileName;
            this.threatName = threatName;
            this.severity = severity;
            this.description = description;
            this.timestamp = timestamp;
        }

        public String getFileName() { return fileName; }
        public String getThreatName() { return threatName; }
        public String getSeverity() { return severity; }
        public String getDescription() { return description; }
        public LocalDateTime getTimestamp() { return timestamp; }
    }
}