import mqtt from 'mqtt';
import fs from 'fs';
import path from 'path';
import { zip, getUUID } from '../../utils.js';

export default class MqttFilePublisher {
    constructor(config) {
        this.config = config;
        this.client = null;
        this.sessionUuid = config.sessionUuid;
        
        if (!config.dirPath) {
            throw new Error('dirPath must be specified in config');
        }
        this.dirPath = config.dirPath;

        // Define binary file extensions
        this.binaryExtensions = new Set([
            '.zip', '.png', '.jpg', '.jpeg', '.gif', 
            '.pdf', '.bin', '.exe', '.tar', '.gz'
        ]);

        // Define text file extensions to zip
        this.textExtensions = new Set([
            '.txt', '.json'
        ]);

        // Track file transfer acknowledgments
        this.pendingFiles = new Map();
        this.failedFiles = new Set();
        
        // Retry configuration
        this.maxRetries = config.maxRetries || 3;
        this.retryDelay = config.retryDelay || 1000; // 1 second
    }

    isBinaryFile(filename) {
        const ext = path.extname(filename).toLowerCase();
        return this.binaryExtensions.has(ext);
    }

    isTextFile(filename) {
        const ext = path.extname(filename).toLowerCase();
        return this.textExtensions.has(ext);
    }

    async createTextFilesZip(files) {
        const textFiles = files.filter(file => this.isTextFile(file));
        if (textFiles.length === 0) return;

        const zipFileName = `TXT-TO-UNZIP-${getUUID()}.zip`;
        const zipFilePath = path.join(this.dirPath, zipFileName);
        const tempDir = path.join(this.dirPath, '_temp_zip_dir');
        
        // Create temporary directory
        if (!fs.existsSync(tempDir)) {
            fs.mkdirSync(tempDir);
        }

        // Move text files to temp directory
        for (const file of textFiles) {
            const sourcePath = path.join(this.dirPath, file);
            const destPath = path.join(tempDir, file);
            fs.renameSync(sourcePath, destPath);
        }

        // Zip the directory containing all text files
        await zip(tempDir, zipFilePath);

        // Clean up temp directory
        fs.rmSync(tempDir, { recursive: true, force: true });
    }

    async publishFile(file, retryCount = 0) {
        const filePath = path.join(this.dirPath, file);
        const isBinary = this.isBinaryFile(file);
        const fileContent = fs.readFileSync(filePath);
        const fileStats = fs.statSync(filePath);
        
        console.log(`Processing ${file} (${(fileStats.size / 1024 / 1024).toFixed(2)} MB) - ${isBinary ? 'binary' : 'text'} file${retryCount > 0 ? ` (Retry ${retryCount}/${this.maxRetries})` : ''}`);
        
        // Track this file
        this.pendingFiles.set(file, true);

        // Create message payload
        const payload = {
            sessionUuid: this.sessionUuid,
            fileName: file,
            fileType: isBinary ? 'binary' : 'text',
            size: fileStats.size,
            content: fileContent.toString('base64')
        };

        try {
            // Publish file and wait for acknowledgment
            await new Promise((resolveFile, rejectFile) => {
                const fileTimeout = setTimeout(() => {
                    console.error(`Timeout waiting for acknowledgment of file ${file}`);
                    rejectFile(new Error('Timeout'));
                }, 30000); // 30 second timeout per file

                this.client.publish('fileTransfer', JSON.stringify(payload), {}, (err) => {
                    if (err) {
                        clearTimeout(fileTimeout);
                        rejectFile(err);
                    }
                });

                // Wait for acknowledgment through message handler
                const checkInterval = setInterval(() => {
                    if (!this.pendingFiles.has(file)) {
                        clearInterval(checkInterval);
                        clearTimeout(fileTimeout);
                        if (this.failedFiles.has(file)) {
                            rejectFile(new Error('File transfer failed'));
                        } else {
                            resolveFile();
                        }
                    }
                }, 100);
            });
        } catch (error) {
            // Handle retry logic
            if (retryCount < this.maxRetries) {
                console.log(`Retrying file ${file} after ${this.retryDelay}ms...`);
                await new Promise(resolve => setTimeout(resolve, this.retryDelay));
                return this.publishFile(file, retryCount + 1);
            } else {
                console.error(`Failed to transfer file ${file} after ${this.maxRetries} retries`);
                this.failedFiles.add(file);
            }
        }
    }

    async publish(reportFilePath) {
        console.log("Running MqttFilePublisher");
        
        return new Promise((resolve, reject) => {
            try {
                // Connect to the MQTT broker
                this.client = mqtt.connect(this.config.loadOrchestrator);

                this.client.on('connect', async () => {
                    console.log('Connected to MQTT broker for publishing results');
                    
                    try {
                        // Subscribe to acknowledgment topics
                        this.client.subscribe(`fileTransfer/${this.sessionUuid}/init`);
                        this.client.subscribe(`fileTransfer/${this.sessionUuid}/ack`);

                        // Handle acknowledgments
                        this.client.on('message', (topic, message) => {
                            const response = JSON.parse(message.toString());
                            
                            if (topic === `fileTransfer/${this.sessionUuid}/init`) {
                                if (response.status === 'ready') {
                                    console.log(`Load orchestrator ready to receive ${response.expectedFiles} files`);
                                }
                            } else if (topic === `fileTransfer/${this.sessionUuid}/ack`) {
                                if (response.status === 'success') {
                                    console.log(`File ${response.fileName} successfully received by load orchestrator`);
                                    this.failedFiles.delete(response.fileName);
                                    this.pendingFiles.delete(response.fileName);
                                } else if (response.status === 'error') {
                                    console.error(`Error receiving file ${response.fileName}: ${response.error}`);
                                    this.failedFiles.add(response.fileName);
                                    this.pendingFiles.delete(response.fileName);
                                }
                            }
                        });

                        // Check if directory exists
                        if (!fs.existsSync(this.dirPath)) {
                            throw new Error(`Directory ${this.dirPath} does not exist`);
                        }

                        // Get initial list of files
                        let allFiles = fs.readdirSync(this.dirPath);
                        
                        // Create zip file for text files
                        await this.createTextFilesZip(allFiles);
                        
                        // Get updated list of files (now includes zip file and excludes moved text files)
                        allFiles = fs.readdirSync(this.dirPath);
                        const fileCount = allFiles.length;

                        console.log(`Publishing ${fileCount} files from directory: ${this.dirPath}`);

                        // Send initialization message and wait for acknowledgment
                        await new Promise((resolveInit, rejectInit) => {
                            const initTimeout = setTimeout(() => {
                                rejectInit(new Error('Timeout waiting for load orchestrator initialization'));
                            }, 10000); // 10 second timeout

                            this.client.publish('fileTransferInit', JSON.stringify({
                                sessionUuid: this.sessionUuid,
                                fileCount: fileCount
                            }), {}, (err) => {
                                if (err) {
                                    clearTimeout(initTimeout);
                                    console.error('Error publishing file count:', err);
                                    rejectInit(err);
                                } else {
                                    console.log(`Published file count: ${fileCount}`);
                                    clearTimeout(initTimeout);
                                    resolveInit();
                                }
                            });
                        });

                        // Send each file with retry support
                        for (const file of allFiles) {
                            await this.publishFile(file);
                            // Add a small delay between files to prevent overwhelming the broker
                            await new Promise(resolve => setTimeout(resolve, 100));
                        }

                        // Check if any files failed after all retries
                        if (this.failedFiles.size > 0) {
                            const error = new Error(`Failed to transfer files after all retries: ${Array.from(this.failedFiles).join(', ')}`);
                            this.client.end();
                            reject(error);
                            return;
                        }

                        console.log('All files published and verified successfully');
                        this.client.end();
                        resolve();
                    } catch (error) {
                        console.error('Error during file publishing:', error);
                        this.client.end();
                        reject(error);
                    }
                });

                this.client.on('error', (err) => {
                    console.error('MQTT client error:', err);
                    reject(err);
                });

            } catch (error) {
                console.error('Error in publish:', error);
                reject(error);
            }
        });
    }
} 