import Aedes from 'aedes';
import { createServer } from 'net';
import mqtt from 'mqtt';
import fs from 'fs';
import path from 'path';
import { fileURLToPath } from 'url';
import allure from 'allure-commandline';
import rimraf from 'rimraf';
import VaultService from '../services/vault-service.js';
import ObjectStorePublisher from '../reporter/publishers/objectstore-publisher.js';
import tools from '@eui/tools/scripts/utils/tools.js';
import { unzip } from '../utils.js';

const { vaultUser, vaultPassword } = tools.getArgs();
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

(async function () {
    let aedes = new Aedes();
    let server = createServer(aedes.handle);
    let port = 1883;

    // Create directories for storing received files
    const filesBaseDir = path.join(__dirname, 'received_files');
    if (!fs.existsSync(filesBaseDir)) {
        fs.mkdirSync(filesBaseDir);
    }

    // Track test sessions and their files
    let testSessions = {};
    let fileTransfers = new Map();

    function cleanupSessionDirectory(sessionUuid) {
        const sessionDir = path.join(filesBaseDir, sessionUuid);
        if (fs.existsSync(sessionDir)) {
            console.log(`Cleaning up existing session directory: ${sessionDir}`);
            rimraf.sync(sessionDir);
        }
        fs.mkdirSync(sessionDir);
        console.log(`Created clean session directory: ${sessionDir}`);
    }

    server.listen(port, function () {
        console.log('server listening on port', port);
    });

    // Create MQTT client with explicit port
    let client = mqtt.connect('mqtt://127.0.0.1:1883', {
        clientId: 'load-orchestrator-' + Math.random().toString(16).substring(2, 8)
    });

    // Topic handler functions
    const topicHandlers = {
        'fileTransferInit': handleFileTransferInit,
        'fileTransfer': handleFileTransfer,
        'clientRegistrationQ': handleClientRegistration,
        'workerRegistrationQ': handleWorkerRegistration,
        'updateState': handleUpdateState,
        'syncStatus': () => printSyncStatusSummary(),
        'checkWorkerStart': handleCheckWorkerStart,

        'workerCrash': handleWorkerCrash,
        'error': handleError
    };

    function handleFileTransferInit(message) {
        const data = JSON.parse(message);
        const { sessionUuid, fileCount } = data;

        // Create session directory if it doesn't exist
        const sessionDir = path.join(filesBaseDir, sessionUuid);
        if (!fs.existsSync(sessionDir)) {
            fs.mkdirSync(sessionDir);
        }

        // Initialize or update session tracking
        if (!fileTransfers.has(sessionUuid)) {
            fileTransfers.set(sessionUuid, {
                clientsExpectedFiles: [],
                receivedFiles: 0
            });
        }

        // Add this client's expected file count
        fileTransfers.get(sessionUuid).clientsExpectedFiles.push(fileCount);
        
        console.log(`Initialized file transfer for session ${sessionUuid}, expecting ${fileCount} files`);
        
        // Send acknowledgment
        client.publish(`fileTransfer/${sessionUuid}/init`, JSON.stringify({
            status: 'ready',
            expectedFiles: fileCount
        }), (err) => {
            if (err) console.error('Error publishing file transfer init acknowledgment:', err);
        });
    }

    function handleFileTransfer(message) {
        const data = JSON.parse(message);
        const { sessionUuid, fileName, content } = data;

        if (!fileTransfers.has(sessionUuid)) {
            console.error(`Received file for unknown session: ${sessionUuid}`);
            return;
        }

        const sessionDir = path.join(filesBaseDir, sessionUuid);
        const filePath = path.join(sessionDir, fileName);
        const fileContent = Buffer.from(content, 'base64');

        try {
            fs.writeFileSync(filePath, fileContent);
            const sessionData = fileTransfers.get(sessionUuid);
            sessionData.receivedFiles++;
            
            console.log(`Received and saved file: ${fileName} for session ${sessionUuid}`);
            
            // Send acknowledgment
            client.publish(`fileTransfer/${sessionUuid}/ack`, JSON.stringify({
                status: 'success',
                fileName: fileName
            }), (err) => {
                if (err) console.error('Error publishing file transfer success acknowledgment:', err);
            });

            // Check if session is complete
            checkSessionComplete(sessionUuid);

        } catch (error) {
            console.error(`Error saving file ${fileName}:`, error);
            // Send error notification
            client.publish(`fileTransfer/${sessionUuid}/ack`, JSON.stringify({
                status: 'error',
                fileName: fileName,
                error: error.message
            }), (err) => {
                if (err) console.error('Error publishing file transfer error acknowledgment:', err);
            });
        }
    }

    function startOrResumeWorkerRampup(sessionUuid) {
        const session = testSessions[sessionUuid];
        const workerRampupDelay = session.workerRampupDelay;
        
        if (!workerRampupDelay || workerRampupDelay <= 0) {
            console.log(`[Rampup] No worker rampup delay configured for session ${sessionUuid}`);
            return;
        }
        
        if (!session.rampupState) {
            console.log(`[Rampup] No rampup state found for session ${sessionUuid}`);
            return;
        }
        
        if (session.rampupState.isRampingUp) {
            console.log(`[Rampup] Session ${sessionUuid} is already ramping up`);
            return;
        }
        
        session.rampupState.isRampingUp = true;
        session.rampupState.isPaused = false;
        
        if (!session.rampupState.rampupStarted) {
            console.log(`[Rampup] Starting worker rampup for session ${sessionUuid} with ${workerRampupDelay}ms delay per worker`);
            session.rampupState.rampupStarted = true;
            session.rampupState.rampupStartTime = Date.now();
        } else {
            console.log(`[Rampup] Resuming worker rampup for session ${sessionUuid} from global worker index ${session.rampupState.nextGlobalWorkerIndex}`);
        }
        
        scheduleNextWorker(sessionUuid);
    }
    
    function scheduleNextWorker(sessionUuid) {
        const session = testSessions[sessionUuid];
        
        if (!session.rampupState.isRampingUp) {
            console.log(`[Rampup] Rampup stopped for session ${sessionUuid}`);
            return;
        }
        
        // Find all available workers that should have started by now
        const availableWorkers = findAllAvailableWorkers(sessionUuid);
        const currentTime = Date.now();
        const rampupStartTime = session.rampupState.rampupStartTime || currentTime;
        
        if (availableWorkers.length === 0) {
            // No more workers available - pause rampup
            console.log(`[Rampup] No more workers available for session ${sessionUuid} at global index ${session.rampupState.nextGlobalWorkerIndex} - pausing rampup`);
            session.rampupState.isRampingUp = false;
            session.rampupState.isPaused = true;
            return;
        }
        
        // Schedule only the next worker in sequence (maintain proper order)
        if (availableWorkers.length > 0) {
            const worker = availableWorkers[0]; // Take the first available worker (lowest global index)
            const { clientUuid, localWorkerIndex, globalWorkerIndex, expectedStartTime, isCatchUp } = worker;
            
            // Calculate delay ensuring minimum workerRampupDelay between any two worker starts
            let delay = 0;
            const now = currentTime;
            const lastWorkerStartTime = session.rampupState.lastWorkerStartTime;
            
            if (lastWorkerStartTime) {
                // Ensure at least workerRampupDelay has passed since the last worker started
                const timeSinceLastWorker = now - lastWorkerStartTime;
                const minimumDelay = session.workerRampupDelay - timeSinceLastWorker;
                delay = Math.max(0, minimumDelay);
                
                if (isCatchUp && delay > 0) {
                    console.log(`[Rampup] Catch-up worker ${localWorkerIndex} on client ${clientUuid} must wait ${delay}ms to maintain minimum ${session.workerRampupDelay}ms interval`);
                }
            } else {
                // First worker can start immediately
                delay = 0;
            }
            
            // For non-catch-up workers, also respect the original expected timing
            if (!isCatchUp) {
                const expectedDelay = Math.max(0, expectedStartTime - now);
                delay = Math.max(delay, expectedDelay);
            }
            
            console.log(`[Rampup] Scheduling worker ${localWorkerIndex} on client ${clientUuid} (global worker ${globalWorkerIndex + 1}) after ${delay}ms ${isCatchUp ? '(catch-up)' : ''}`);
            
            setTimeout(() => {
                // Check if rampup is still active (could have been stopped)
                if (!session.rampupState.isRampingUp) {
                    console.log(`[Rampup] Rampup was stopped for session ${sessionUuid} - skipping worker ${localWorkerIndex} on ${clientUuid}`);
                    return;
                }
                
                const actualStartTime = Date.now();
                console.log(`[Worker Rampup] Starting worker ${localWorkerIndex} on client ${clientUuid} (global worker ${globalWorkerIndex + 1})`);
                
                // Store worker start permission in session state for workers to check
                if (!testSessions[sessionUuid].workerStartPermissions) {
                    testSessions[sessionUuid].workerStartPermissions = {};
                }
                if (!testSessions[sessionUuid].workerStartPermissions[clientUuid]) {
                    testSessions[sessionUuid].workerStartPermissions[clientUuid] = {};
                }
                
                testSessions[sessionUuid].workerStartPermissions[clientUuid][localWorkerIndex] = {
                    canStart: true,
                    startTime: actualStartTime,
                    rampupDelay: session.workerRampupDelay,
                    globalIndex: globalWorkerIndex
                };
                
                console.log(`[Worker Rampup] Stored start permission for worker ${localWorkerIndex} on client ${clientUuid} - worker will poll for this permission`);
                
                // No longer sending MQTT messages - workers will poll for permissions instead
                // This ensures no race conditions with message timing
                
                // Continue with the scheduling logic (no network call needed)
                console.log(`[Rampup] Worker ${localWorkerIndex} on ${clientUuid} (global ${globalWorkerIndex + 1}) permission granted successfully`);
                
                // Record the actual start time
                session.rampupState.lastWorkerStartTime = actualStartTime;
                
                // Move to next worker
                session.rampupState.nextGlobalWorkerIndex = globalWorkerIndex + 1;
                
                // Check if we've ramped up all workers
                if (session.rampupState.nextGlobalWorkerIndex >= session.users) {
                    console.log(`[Rampup] All ${session.users} workers have been given permission for session ${sessionUuid} - sending global start signal`);
                    session.rampupState.isRampingUp = false;
                    
                    // Send global start signal since all workers are now ramped up
                    client.publish(sessionUuid + "/all", "start", (err) => {
                        if (err) console.error('Error publishing start message after rampup completion:', err);
                        else console.log(`Global start signal sent successfully to ${sessionUuid}/all after rampup completion`);
                    });
                } else {
                    // Schedule the next worker immediately (delay will be calculated based on lastWorkerStartTime)
                    setTimeout(() => {
                        if (session.rampupState.isRampingUp) {
                            scheduleNextWorker(sessionUuid);
                        }
                    }, 50); // Small delay before checking for next worker
                }
            }, delay);
        }
    }
    
    function findAllAvailableWorkers(sessionUuid) {
        const session = testSessions[sessionUuid];
        const currentTime = Date.now();
        const rampupStartTime = session.rampupState.rampupStartTime || currentTime;
        const availableWorkers = [];
        let currentGlobalIndex = 0;
        
        // Find the next worker in sequence that should be scheduled
        const targetGlobalIndex = session.rampupState.nextGlobalWorkerIndex;
        
        // Iterate through all clients and their workers to find the target worker
        for (const [clientUuid, clientData] of Object.entries(session.clients)) {
            for (let localWorkerIndex = 0; localWorkerIndex < clientData.workerCount; localWorkerIndex++) {
                if (currentGlobalIndex === targetGlobalIndex && currentGlobalIndex < session.users) {
                    const expectedStartTime = rampupStartTime + (currentGlobalIndex * session.workerRampupDelay);
                    let isCatchUp = false;
                    
                    // Check if this is the first worker from this client in the global sequence
                    const isFirstWorkerFromThisClient = localWorkerIndex === 0;
                    
                    if (isFirstWorkerFromThisClient && expectedStartTime <= currentTime) {
                        // Catch-up logic: first worker of a late client can start immediately
                        isCatchUp = true;
                        console.log(`[Rampup] First worker (${localWorkerIndex}) of late client ${clientUuid} gets catch-up (was due at ${new Date(expectedStartTime).toISOString()})`);
                    }
                    
                    availableWorkers.push({
                        clientUuid,
                        localWorkerIndex,
                        globalWorkerIndex: currentGlobalIndex,
                        expectedStartTime,
                        isCatchUp
                    });
                    
                    // Only return one worker at a time to maintain sequence
                    break;
                }
                
                currentGlobalIndex++;
            }
            
            // If we found the target worker, stop looking
            if (availableWorkers.length > 0) {
                break;
            }
        }
        
        console.log(`[Rampup] Looking for worker at global index ${targetGlobalIndex}. Found: ${availableWorkers.length > 0 ? 'Yes' : 'No'}. Total registered workers: ${currentGlobalIndex}, target users: ${session.users}`);
        
        return availableWorkers;
    }

    async function handleClientRegistration(message) {
        let sessionConfig = JSON.parse(message);
        let sessionUuid = sessionConfig.sessionUuid;
        let sessionclient = sessionConfig.client;
        let clientUuid = sessionclient.uuid;
        delete sessionclient.uuid;
        delete sessionConfig.sessionUuid;
        delete sessionConfig.client;

        if (sessionConfig.resetSession && sessionUuid in testSessions) {
            delete testSessions[sessionUuid];
            // Reset file transfers for this session
            if (fileTransfers.has(sessionUuid)) {
                fileTransfers.delete(sessionUuid);
            }
            cleanupSessionDirectory(sessionUuid);
        }

        if (sessionUuid in testSessions) {
            if (clientUuid in testSessions[sessionUuid].clients) {
                console.log(`Client already registered: ${sessionUuid} / ${clientUuid}`);
            } else {
                let workerStart = testSessions[sessionUuid].workers + 1;
                let workerEnd = testSessions[sessionUuid].workers + sessionclient.workerCount;
                testSessions[sessionUuid].clients[clientUuid] = sessionclient;
                testSessions[sessionUuid].workers = workerEnd;
                client.publish(sessionUuid + "/" + clientUuid,
                    JSON.stringify({ "workerStart": workerStart, "workerEnd": workerEnd }), 
                    (err) => {
                        if (err) console.error('Error publishing worker range:', err);
                    });
                console.log(`New Client added: ${sessionUuid} / ${clientUuid}`);
                
                // If worker rampup is configured, start or resume the rampup process
                const workerRampupDelay = testSessions[sessionUuid].workerRampupDelay;
                if (workerRampupDelay && workerRampupDelay > 0) {
                    const rampupState = testSessions[sessionUuid].rampupState;
                    if (rampupState && (rampupState.isPaused || !rampupState.rampupStarted)) {
                        console.log(`[Rampup] New client added - attempting to start/resume rampup for session ${sessionUuid}`);
                        startOrResumeWorkerRampup(sessionUuid);
                    }
                }
            }
        } else {
            // Clean up session directory when new session starts
            cleanupSessionDirectory(sessionUuid);
            
            sessionConfig.clients = {};
            sessionConfig.clients[clientUuid] = sessionclient;
            testSessions[sessionUuid] = sessionConfig;
            testSessions[sessionUuid].workers = sessionclient.workerCount;
            
            // Initialize rampup state tracking
            if (sessionConfig.workerRampupDelay && sessionConfig.workerRampupDelay > 0) {
                testSessions[sessionUuid].rampupState = {
                    nextGlobalWorkerIndex: 0,
                    isRampingUp: false,
                    isPaused: false,
                    rampupStarted: false,
                    lastWorkerStartTime: null // Track when the last worker actually started
                };
            }
            
            client.publish(sessionUuid + "/" + clientUuid,
                JSON.stringify({ "workerStart": 1, "workerEnd": sessionclient.workerCount }), 
                (err) => {
                    if (err) console.error('Error publishing initial worker range:', err);
                });
            console.log("NEW Session created: " + sessionUuid);
            console.log("First Client added with uuid: " + clientUuid);
            
            // If worker rampup is configured, start the rampup process immediately
            const workerRampupDelay = testSessions[sessionUuid].workerRampupDelay;
            if (workerRampupDelay && workerRampupDelay > 0) {
                startOrResumeWorkerRampup(sessionUuid);
            }
        }

        // Check if we need to send global start signal (only for non-rampup scenarios)
        const workerRampupDelay = testSessions[sessionUuid].workerRampupDelay;
        
        if (testSessions[sessionUuid].workers >= testSessions[sessionUuid].users) {
            console.log("Workers goal reached for test: " + sessionUuid + " CURRENT SETUP: " + JSON.stringify(testSessions[sessionUuid]));
            
            if (!workerRampupDelay || workerRampupDelay <= 0) {
                // No worker rampup delay - send immediate global start signal
                console.log(`No workerRampupDelay configured - sending immediate global start signal`);
                client.publish(sessionUuid + "/all", "start", 
                    (err) => {
                        if (err) console.error('Error publishing start message:', err);
                        else console.log(`Global start signal sent successfully to ${sessionUuid}/all`);
                    });
            } else {
                console.log(`Worker rampup is active - global start signal will be sent when all workers are ramped up`);
            }
            
        } else {
            console.log("Workers goal NOT YET reached for test: " + sessionUuid + " CURRENT SETUP: " + JSON.stringify(testSessions[sessionUuid]));
        }
    }

    function handleWorkerRegistration(message) {
        console.log("Received workers");
        let workerConfig = JSON.parse(message);
        console.log(workerConfig);
        let sessionUuid = workerConfig.sessionUuid;
        let clientUuid = workerConfig.clientUuid;

        if (sessionUuid in testSessions) {
            if (clientUuid in testSessions[sessionUuid].clients) {
                testSessions[sessionUuid].clients[clientUuid].workers = workerConfig.workers;
                console.log(`workerRegistrationQ Client workers added for: ${sessionUuid} / ${clientUuid}`);
                
                // If worker rampup is paused (waiting for workers), try to resume it
                const session = testSessions[sessionUuid];
                if (session.rampupState && session.rampupState.isPaused && !session.rampupState.isRampingUp) {
                    console.log(`[Rampup] Attempting to resume paused rampup for session ${sessionUuid} after worker registration`);
                    startOrResumeWorkerRampup(sessionUuid);
                }
            } else {
                console.log(`workerRegistrationQ Client unknown: ${sessionUuid} / ${clientUuid}`);
            }
        } else {
            console.log("workerRegistrationQ Session unknown: " + sessionUuid);
        }
    }

    function handleCheckWorkerStart(message) {
        try {
            const request = JSON.parse(message);
            const { sessionUuid, clientUuid, workerIndex, replyTopic } = request;
            
            console.log(`[CheckWorkerStart] Checking start permission for worker ${workerIndex} on client ${clientUuid}`);
            
            if (sessionUuid in testSessions) {
                const permissions = testSessions[sessionUuid].workerStartPermissions;
                if (permissions && permissions[clientUuid] && permissions[clientUuid][workerIndex]) {
                    const permission = permissions[clientUuid][workerIndex];
                    if (permission.canStart) {
                        console.log(`[CheckWorkerStart] Worker ${workerIndex} on ${clientUuid} has start permission`);
                        
                        // Calculate the exact time this worker should start
                        const session = testSessions[sessionUuid];
                        const rampupStartTime = session.rampupState.rampupStartTime;
                        const targetStartTime = rampupStartTime + (permission.globalIndex * permission.rampupDelay);
                        
                        client.publish(replyTopic, JSON.stringify({
                            canStart: true,
                            startTime: permission.startTime,
                            rampupDelay: permission.rampupDelay,
                            globalIndex: permission.globalIndex,
                            targetStartTime: targetStartTime
                        }));
                        return;
                    }
                }
            }
            
            console.log(`[CheckWorkerStart] Worker ${workerIndex} on ${clientUuid} does NOT have start permission yet`);
            client.publish(replyTopic, JSON.stringify({
                canStart: false
            }));
            
        } catch (error) {
            console.error('Error handling worker start check:', error);
        }
    }

    function handleUpdateState(message) {
        let stateConfig = JSON.parse(message);
        console.log(stateConfig);
        let sessionUuid = stateConfig.sessionUuid;
        let clientUuid = stateConfig.clientUuid;

        if (sessionUuid in testSessions) {
            if (clientUuid in testSessions[sessionUuid].clients) {
                if ("worker" in stateConfig) {
                    if (testSessions[sessionUuid].clients[clientUuid].workers[stateConfig.worker]) {
                        testSessions[sessionUuid].clients[clientUuid].workers[stateConfig.worker].state = stateConfig.state;
                        testSessions[sessionUuid].clients[clientUuid].workers[stateConfig.worker].syncpoint = stateConfig.syncpoint;
                        
                        // Get the list of workers that are NOT permanently failed
                        const activeWorkers = getActiveWorkers(sessionUuid, clientUuid);
                        const workersInTargetState = testSessions[sessionUuid].clients[clientUuid].workers.filter(
                            (item, index) => {
                                // Only count workers that are active (not permanently failed) and in the target state
                                return activeWorkers.includes(index) && 
                                       item.state == stateConfig.state && 
                                       item.syncpoint == stateConfig.syncpoint;
                            }
                        );
                        
                        // Enhanced logging with more details
                        const totalWorkers = testSessions[sessionUuid].clients[clientUuid].workers.length;
                        const failedWorkerCount = totalWorkers - activeWorkers.length;
                        
                        console.log(`[SYNC STATUS] Client ${clientUuid}: ${workersInTargetState.length}/${activeWorkers.length} active workers reached state "${stateConfig.state}" syncpoint ${stateConfig.syncpoint} (${failedWorkerCount} failed workers excluded)`);
                        
                        // Check if all ACTIVE workers have reached the target state
                        if (workersInTargetState.length === activeWorkers.length) {
                            testSessions[sessionUuid].clients[clientUuid].state = stateConfig.state;
                            testSessions[sessionUuid].clients[clientUuid].syncpoint = stateConfig.syncpoint;
                            console.log('\x1b[32m%s\x1b[0m', `[CLIENT SYNC] ✓ All ${activeWorkers.length} active workers on ${clientUuid} reached state "${stateConfig.state}" syncpoint ${stateConfig.syncpoint}`);
                            setTimeout(() => {
                                client.publish(sessionUuid + "/" + clientUuid, JSON.stringify({ "state": stateConfig.state }), (err) => {
                                    if (err) console.error('Error publishing worker state:', err);
                                });
                            }, 5000); // 5 second delay
                            if (stateConfig.syncpoint) {
                                let clientsInSync = true;
                                let totalActiveClients = 0;
                                let clientsAtSyncPoint = 0;
                                
                                Object.keys(testSessions[sessionUuid].clients).forEach(checkClientUuid => {
                                    const clientActiveWorkers = getActiveWorkers(sessionUuid, checkClientUuid);
                                    // Only count clients that have active workers
                                    if (clientActiveWorkers.length > 0) {
                                        totalActiveClients++;
                                        if (testSessions[sessionUuid].clients[checkClientUuid].syncpoint === stateConfig.syncpoint) {
                                            clientsAtSyncPoint++;
                                        } else {
                                            clientsInSync = false;
                                        }
                                    }
                                });
                                
                                console.log(`[SESSION SYNC] ${clientsAtSyncPoint}/${totalActiveClients} clients at syncpoint ${stateConfig.syncpoint} in session ${sessionUuid}`);
                                
                                if (clientsInSync && totalActiveClients > 0) {
                                    console.log('\x1b[36m%s\x1b[0m', `[GLOBAL SYNC] ✓ All ${totalActiveClients} clients synchronized at syncpoint ${stateConfig.syncpoint} - broadcasting sync signal to session ${sessionUuid}`);
                                    setTimeout(() => {
                                        client.publish(sessionUuid + "/all", 'sync', (err) => {
                                            if (err) console.error('Error publishing sync message:', err);
                                        });
                                    }, 5000); // 5 second delay
                                }
                            }
                        }
                    } else {
                        console.log(`updateState worker ${stateConfig.worker} unknown: ${sessionUuid} / ${clientUuid}`);
                    }
                } else {
                    testSessions[sessionUuid].clients[clientUuid].state = stateConfig.state;
                    console.log(`State ${stateConfig.state} received for ${sessionUuid} / ${clientUuid}`);
                    setTimeout(() => {
                        client.publish(sessionUuid + "/" + clientUuid, JSON.stringify({ "state": stateConfig.state }), (err) => {
                            if (err) console.error('Error publishing client state:', err);
                        });
                    }, 1000); // 1 second delay
                }
            } else {
                console.log(`updateState Client unknown: ${sessionUuid} / ${clientUuid}`);
            }
        } else {
            console.log("updateState Session unknown: " + sessionUuid);
        }
    }

    // Helper function to get list of active (non-failed) worker indices for a client
    function getActiveWorkers(sessionUuid, clientUuid) {
        const totalWorkerCount = testSessions[sessionUuid].clients[clientUuid].workerCount || 0;
        const allWorkerIndices = Array.from({length: totalWorkerCount}, (_, i) => i);
        
        // Get permanently failed workers for this client
        const failedWorkers = testSessions[sessionUuid].failedWorkers?.[clientUuid] || new Set();
        
        // Return workers that are not permanently failed
        const activeWorkers = allWorkerIndices.filter(index => !failedWorkers.has(index));
        
        return activeWorkers;
    }

    // Helper function to print comprehensive sync status summary
    function printSyncStatusSummary() {
        console.log('\n' + '='.repeat(80));
        console.log('\x1b[36m%s\x1b[0m', '[SYNC STATUS SUMMARY]');
        console.log('='.repeat(80));
        
        Object.keys(testSessions).forEach(sessionUuid => {
            const session = testSessions[sessionUuid];
            console.log(`\nSession: ${sessionUuid}`);
            console.log(`Error State: ${session.error ? 'YES' : 'NO'}`);
            
            Object.keys(session.clients).forEach(clientUuid => {
                const client = session.clients[clientUuid];
                const activeWorkers = getActiveWorkers(sessionUuid, clientUuid);
                const totalWorkers = client.workers.length;
                const failedWorkers = totalWorkers - activeWorkers.length;
                
                console.log(`  Client ${clientUuid}:`);
                console.log(`    Workers: ${activeWorkers.length} active, ${failedWorkers} failed (${totalWorkers} total)`);
                console.log(`    State: ${client.state || 'N/A'}`);
                console.log(`    Syncpoint: ${client.syncpoint || 'N/A'}`);
                
                // Show individual worker states for debugging
                if (activeWorkers.length > 0 && client.workers[0]) {
                    const workerStates = {};
                    client.workers.forEach((worker, index) => {
                        if (activeWorkers.includes(index)) {
                            const stateKey = `${worker.state || 'unknown'}:${worker.syncpoint || 'none'}`;
                            workerStates[stateKey] = (workerStates[stateKey] || 0) + 1;
                        }
                    });
                    
                    console.log(`    Worker Distribution:`);
                    Object.entries(workerStates).forEach(([state, count]) => {
                        console.log(`      ${count} workers at ${state}`);
                    });
                }
            });
        });
        
        console.log('='.repeat(80) + '\n');
    }



    function handleWorkerCrash(message) {
        try {
            const crashData = JSON.parse(message);
            const { sessionUuid, clientUuid, workerIndex, exitCode, errorType, specs, timestamp } = crashData;
            
            console.log('\x1b[31m%s\x1b[0m', `[ORCHESTRATOR] Worker ${workerIndex} crashed on ${clientUuid} with ${errorType} (exit code: ${exitCode})`);
            
            if (specs && specs.length > 0) {
                console.log('\x1b[33m%s\x1b[0m', `[ORCHESTRATOR] Crashed specs: ${specs.join(', ')}`);
            }
            
            // Track crashes for monitoring and session health assessment
            if (testSessions[sessionUuid] && testSessions[sessionUuid].clients[clientUuid]) {
                if (!testSessions[sessionUuid].workerCrashes) {
                    testSessions[sessionUuid].workerCrashes = {};
                }
                
                if (!testSessions[sessionUuid].workerCrashes[clientUuid]) {
                    testSessions[sessionUuid].workerCrashes[clientUuid] = [];
                }
                
                const crashInfo = {
                    workerIndex,
                    exitCode,
                    errorType,
                    specs,
                    timestamp
                };
                
                testSessions[sessionUuid].workerCrashes[clientUuid].push(crashInfo);
                
                // Immediately mark crashed worker as permanently failed to exclude from sync calculations
                if (!testSessions[sessionUuid].failedWorkers) {
                    testSessions[sessionUuid].failedWorkers = {};
                }
                
                if (!testSessions[sessionUuid].failedWorkers[clientUuid]) {
                    testSessions[sessionUuid].failedWorkers[clientUuid] = new Set();
                }
                
                testSessions[sessionUuid].failedWorkers[clientUuid].add(workerIndex);
                console.log('\x1b[33m%s\x1b[0m', `[ORCHESTRATOR] Worker ${workerIndex} marked as failed - excluded from sync point calculations`);
                
                // Check session health after marking worker as failed
                const totalWorkers = testSessions[sessionUuid].clients[clientUuid].workerCount || 0;
                const failedCount = testSessions[sessionUuid].failedWorkers[clientUuid].size;
                const failureThreshold = Math.ceil(totalWorkers * 0.3);
                
                console.log(`[ORCHESTRATOR] Failure status for ${clientUuid}: ${failedCount}/${totalWorkers} workers failed (threshold: ${failureThreshold})`);
                
                if (failedCount >= failureThreshold) {
                    console.log('\x1b[31m%s\x1b[0m', `[ORCHESTRATOR] Too many workers failed (${failedCount}/${totalWorkers}) for ${sessionUuid}/${clientUuid}, marking session as failed`);
                    testSessions[sessionUuid].error = true;
                    
                    // Stop the session due to excessive failures
                    client.publish(`${sessionUuid}/all`, 'error', {}, (err) => {
                        if (err) {
                            console.error('Error forwarding error message due to worker failures:', err);
                        } else {
                            console.log(`[SESSION ERROR] Session ${sessionUuid} stopped due to excessive worker failures`);
                        }
                    });
                }
                
                // Log detailed crash information
                if (errorType === 'v8_fatal_error') {
                    console.log('\x1b[31m%s\x1b[0m', `[ORCHESTRATOR] V8 Fatal Error confirmed on worker ${workerIndex} - worker excluded from further sync points`);
                } else {
                    console.log('\x1b[33m%s\x1b[0m', `[ORCHESTRATOR] Worker crash logged for session health monitoring`);
                }
            }
            
        } catch (err) {
            console.error('Error handling worker crash message:', err);
        }
    }





    async function checkSessionComplete(sessionUuid) {
        const sessionData = fileTransfers.get(sessionUuid);
        const sessionInfo = testSessions[sessionUuid];
        
        if (!sessionData || !sessionInfo) return false;

        // Check if we've received all expected files from all clients
        const expectedTotalFiles = sessionData.clientsExpectedFiles.reduce((sum, count) => sum + count, 0);
        const actualFiles = fs.readdirSync(path.join(filesBaseDir, sessionUuid)).length;

        console.log(`Session ${sessionUuid} status:
            Expected clients: ${Object.keys(sessionInfo.clients).length}
            Reporting clients: ${sessionData.clientsExpectedFiles.length}
            Expected files: ${expectedTotalFiles}
            Received files: ${actualFiles}
            Error state: ${sessionInfo.error ? 'YES' : 'NO'}`);

        if (sessionData.clientsExpectedFiles.length === Object.keys(sessionInfo.clients).length &&
            actualFiles === expectedTotalFiles) {
            console.log(`Session ${sessionUuid} is complete. Processing files...`);
            
            const sessionDir = path.join(filesBaseDir, sessionUuid);

            // If session is in error state, skip report generation and publishing
            if (sessionInfo.error) {
                console.log('\x1b[31m%s\x1b[0m', `Session ${sessionUuid} is in error state - skipping report generation and publishing`);
                
                // Clean up all files
                if (fs.existsSync(sessionDir)) {
                    fs.rmSync(sessionDir, { recursive: true, force: true });
                    console.log(`Session directory ${sessionDir} cleaned up`);
                }
                
                return true;
            }
            
            // Unzip text files
            try {
                const files = fs.readdirSync(sessionDir);
                const textZipFiles = files.filter(file => file.startsWith('TXT-TO-UNZIP-') && file.endsWith('.zip'));
                
                for (const zipFile of textZipFiles) {
                    console.log(`Processing ${zipFile}...`);
                    const tempDir = path.join(sessionDir, '_temp_unzip');
                    if (!fs.existsSync(tempDir)) {
                        fs.mkdirSync(tempDir);
                    }

                    // Unzip to temp directory
                    await unzip(sessionDir, zipFile);

                    // Move all files from temp directory to session directory
                    const extractedFiles = fs.readdirSync(tempDir);
                    for (const file of extractedFiles) {
                        const sourcePath = path.join(tempDir, file);
                        const destPath = path.join(sessionDir, file);
                        fs.renameSync(sourcePath, destPath);
                    }

                    // Clean up
                    fs.rmSync(tempDir, { recursive: true, force: true });
                    fs.unlinkSync(path.join(sessionDir, zipFile));
                }
            } catch (error) {
                console.error('Error unzipping text files:', error);
                return false;
            }

            console.log('Generating Allure report...');
            const reportDir = path.join(sessionDir, 'allure-report');
            const reportCommand = allure(['generate', sessionDir, '--clean', '-o', reportDir]);
            
            try {
                await new Promise((resolve, reject) => {
                    reportCommand.on('exit', code => {
                        if (code === 0) {
                            resolve();
                        } else {
                            reject(new Error(`Allure report generation failed with code ${code}`));
                        }
                    });
                });

                console.log('Allure report generated successfully');

                // Clean up all result files since report is generated
                const files = fs.readdirSync(sessionDir);
                files.forEach(file => {
                    const filePath = path.join(sessionDir, file);
                    // Skip the allure-report directory
                    if (filePath !== reportDir) {
                        if (fs.statSync(filePath).isDirectory()) {
                            fs.rmSync(filePath, { recursive: true, force: true });
                        } else {
                            fs.unlinkSync(filePath);
                        }
                    }
                });
                console.log('Result files cleaned up successfully');

                // Publish to object store if configured
                if (sessionInfo.objectstore && sessionInfo.vaultConfig) {
                    try {
                        // Initialize vault service and get S3 credentials
                        console.log('Initializing vault service and getting S3 credentials...');
                        const vaultService = new VaultService({
                            vaultConfig: {
                                ...sessionInfo.vaultConfig,
                                user: vaultUser,
                                password: vaultPassword
                            }
                        });

                        if (sessionInfo.vaultEnvVars?.s3) {
                            const s3Secrets = await vaultService.getSecret(
                                sessionInfo.vaultEnvVars.s3.path,
                                sessionInfo.vaultEnvVars.s3.key
                            );
                            
                            if (s3Secrets[sessionInfo.vaultEnvVars.s3.key]) {
                                Object.entries(s3Secrets[sessionInfo.vaultEnvVars.s3.key]).forEach(([key, value]) => {
                                    process.env[key] = value;
                                });
                                console.log('S3 credentials obtained successfully');
                            }
                        }

                        // Initialize ObjectStorePublisher
                        console.log('Initializing ObjectStore publisher...');
                        const objectStorePublisher = new ObjectStorePublisher({
                            ...sessionInfo.objectstore,
                            dirPath: reportDir,
                            s3rootKey: sessionInfo.objectstore.prefixRootKeyWithTime ? 
                                `${new Date().toISOString().replace(/[:.]/g, '-')}${sessionInfo.objectstore.s3rootKey || ''}` : 
                                sessionInfo.objectstore.s3rootKey
                        });

                        // Publish the report
                        await objectStorePublisher.publish();
                        console.log('Report published to object store successfully');

                        // Clean up the entire session directory after successful publish
                        fs.rmSync(sessionDir, { recursive: true, force: true });
                        console.log(`Session directory ${sessionDir} cleaned up successfully`);

                    } catch (error) {
                        console.error('Failed to publish report to object store:', error);
                    }
                }

                return true;
            } catch (error) {
                console.error('Error during report generation:', error);
                return false;
            }
        }

        return false;
    }

    // Handle client connection
    client.on('connect', function () {
        console.log('MQTT client connected');
        
        // Subscribe to standard topics
        Object.keys(topicHandlers).forEach(topic => {
            client.subscribe(topic, (err) => {
                if (err) {
                    console.error('Error subscribing to topic:', topic, err);
                } else {
                    console.log('Subscribed to topic:', topic);
                }
            });
        });
    });

    // Handle client errors
    client.on('error', function (err) {
        console.error('MQTT client error:', err);
    });

    // Handle client reconnection
    client.on('reconnect', function () {
        console.log('MQTT client reconnecting...');
    });

    // Handle client disconnection
    client.on('close', function () {
        console.log('MQTT client disconnected');
    });

    // Handle incoming messages
    client.on('message', function (topic, message) {
        const messageStr = message.toString();
        // console.log(`[DEBUG] Received message - Topic: ${topic}, Message: ${messageStr}`);
        
        const handler = topicHandlers[topic];
        if (handler) {
            console.log(`[DEBUG] Calling handler for topic: ${topic}`);
            handler(messageStr);
        } else {
            console.log(`[DEBUG] No handler found for topic: ${topic}`);
        }
    });

    function handleError(message) {
        try {
            console.log('[DEBUG] Error handler called with message:', message);
            const sessionUuid = message;
            console.log('\x1b[31m%s\x1b[0m', `ERROR MESSAGE RECEIVED for session ${sessionUuid}`); // Red color
            console.log('\x1b[31m%s\x1b[0m', `Stopping all tests for session ${sessionUuid}`); // Red color
            
            // Mark session as in error state
            if (testSessions[sessionUuid]) {
                testSessions[sessionUuid].error = true;
                console.log(`Session ${sessionUuid} marked as in error state`);
            }
            
            // Forward the error to the specific session's all channel
            client.publish(`${sessionUuid}/all`, 'error', {}, (err) => {
                if (err) {
                    console.error('Error forwarding error message:', err);
                } else {
                    console.log(`Error message forwarded to session ${sessionUuid}`);
                }
            });
        } catch (err) {
            console.error('Error handling error message:', err);
        }
    }
})();