/**
 * Migration Script: Populate playbook_solution.correct_action for all alerts
 * This determines the "correct first response" for each alert based on keywords.
 */
const db = require('../src/db');

const ACTION_OPTIONS = ['block', 'isolate', 'lock', 'purge'];
const ACTION_LABELS = {
    block: 'BLOCK IP',
    isolate: 'ISOLATE HOST', 
    lock: 'LOCK ACCOUNT',
    purge: 'PURGE INBOX'
};

// Determine correct action based on alert keywords
const determineCorrectAction = (alert) => {
    const title = (alert.title || '').toLowerCase();
    const type = (alert.alert_type || '').toLowerCase();
    const source = (alert.source || '').toLowerCase();
    const desc = (alert.description || '').toLowerCase(); // If available

    // Combined text for search
    const text = title + ' ' + type + ' ' + source + ' ' + desc;

    // Phishing/Email → Purge Inbox
    if (text.includes('phish') || text.includes('spear') || text.includes('email') || 
        text.includes('bec') || text.includes('mail') || text.includes('suspicious subject') ||
        text.includes('sender') || text.includes('inbox')) {
        return 'purge';
    }
    
    // Malware/Ransomware/Endpoint → Isolate Host
    if (text.includes('malware') || text.includes('ransomware') || text.includes('mimikatz') || 
        text.includes('cobalt') || text.includes('beacon') || text.includes('endpoint') ||
        text.includes('script') || text.includes('powershell') || text.includes('c2') ||
        text.includes('command and control') || text.includes('process') || text.includes('file') ||
        text.includes('execution') || text.includes('download')) {
        return 'isolate';
    }
    
    // Identity/Account/Login → Lock Account
    if (text.includes('login') || text.includes('account') || text.includes('credential') ||
        text.includes('brute') || text.includes('identity') || text.includes('privilege') ||
        text.includes('password') || text.includes('access') || text.includes('failed')) {
        return 'lock';
    }
    
    // Default: Block IP (Network-based, Scanning, Firewall)
    return 'block';
};

async function migrate() {
    try {
        console.log("🚀 Starting Migration: Populating correct_action for all alerts...\n");
        
        const result = await db.query("SELECT id, title, alert_type, source, playbook_solution FROM alerts");
        const alerts = result.rows;
        
        console.log(`Found ${alerts.length} alerts to process.\n`);

        let stats = { block: 0, isolate: 0, lock: 0, purge: 0 };

        for (const alert of alerts) {
            // Parse existing solution if any
            let solution = {};
            if (alert.playbook_solution) {
                if (typeof alert.playbook_solution === 'string') {
                    try { solution = JSON.parse(alert.playbook_solution); } catch(e) {}
                } else {
                    solution = alert.playbook_solution;
                }
            }

            // Determine and set correct action
            const correctAction = determineCorrectAction(alert);
            solution.correct_action = correctAction;
            solution.action_label = ACTION_LABELS[correctAction];
            
            stats[correctAction]++;
            
            // Update DB
            await db.query(
                "UPDATE alerts SET playbook_solution = ? WHERE id = ?",
                [JSON.stringify(solution), alert.id]
            );
        }

        console.log("✅ Migration Complete!\n");
        console.log("📊 Action Distribution:");
        console.log(`   BLOCK IP:      ${stats.block}`);
        console.log(`   ISOLATE HOST:  ${stats.isolate}`);
        console.log(`   LOCK ACCOUNT:  ${stats.lock}`);
        console.log(`   PURGE INBOX:   ${stats.purge}`);
        
        process.exit(0);
    } catch (err) {
        console.error("❌ Migration Failed:", err);
        process.exit(1);
    }
}

migrate();
