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,
        '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);
            });
        }
    }

    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}`);
            }
        } 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;
            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 (testSessions[sessionUuid].workers >= testSessions[sessionUuid].users) {
            console.log("Workers goal reached for test: " + sessionUuid + " CURRENT SETUP: " + JSON.stringify(testSessions[sessionUuid]));
            client.publish(sessionUuid + "/all", "start", 
                (err) => {
                    if (err) console.error('Error publishing start message:', err);
                });
        } 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}`);
            } else {
                console.log(`workerRegistrationQ Client unknown: ${sessionUuid} / ${clientUuid}`);
            }
        } else {
            console.log("workerRegistrationQ Session unknown: " + sessionUuid);
        }
    }

    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;
                        if (testSessions[sessionUuid].clients[clientUuid].workers.filter(item => item.state == stateConfig.state && item.syncpoint == stateConfig.syncpoint).length ==
                            testSessions[sessionUuid].clients[clientUuid].workers.length) {
                            testSessions[sessionUuid].clients[clientUuid].state = stateConfig.state;
                            testSessions[sessionUuid].clients[clientUuid].syncpoint = stateConfig.syncpoint;
                            console.log(`All workers in state ${stateConfig.state} with syncpoint ${stateConfig.syncpoint} for ${sessionUuid} / ${clientUuid}`);
                            client.publish(sessionUuid + "/" + clientUuid, JSON.stringify({ "state": stateConfig.state }), (err) => {
                                if (err) console.error('Error publishing worker state:', err);
                            });
                            if (stateConfig.syncpoint) {
                                let clientsInSync = true;
                                Object.keys(testSessions[sessionUuid].clients).forEach(clientUuid => {
                                    if (testSessions[sessionUuid].clients[clientUuid].syncpoint !== stateConfig.syncpoint) {
                                        clientsInSync = false;
                                    }
                                });
                                if (clientsInSync) {
                                    console.log(`All workers in sync ${stateConfig.syncpoint} for session ${sessionUuid}`);
                                    client.publish(sessionUuid + "/all", 'sync', (err) => {
                                        if (err) console.error('Error publishing sync message:', err);
                                    });
                                }
                            }
                        }
                    } 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}`);
                    client.publish(sessionUuid + "/" + clientUuid, JSON.stringify({ "state": stateConfig.state }), (err) => {
                        if (err) console.error('Error publishing client state:', err);
                    });
                }
            } else {
                console.log(`updateState Client unknown: ${sessionUuid} / ${clientUuid}`);
            }
        } else {
            console.log("updateState Session unknown: " + sessionUuid);
        }
    }

    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);
        }
    }
})();