import { api } from "../../scripts/api.js";
import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker } from './workerUtils.js';
import { TIMEOUTS } from './constants.js';

export function setupInterceptor(extension) {
    api.queuePrompt = async (number, prompt) => {
        if (extension.isEnabled) {
            const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0;
            const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0;
            
            if (hasCollector || hasDistUpscale) {
                const result = await executeParallelDistributed(extension, prompt);
                // Check status after dispatching jobs
                setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
                return result;
            }
        }
        return extension.originalQueuePrompt(number, prompt);
    };
}

export async function executeParallelDistributed(extension, promptWrapper) {
    try {
        const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution
        const enabledWorkers = extension.enabledWorkers;
        
        // Pre-flight health check on all enabled workers
        const activeWorkers = await performPreflightCheck(extension, enabledWorkers);
        
        // Case: Enabled workers but all offline
        if (activeWorkers.length === 0 && enabledWorkers.length > 0) {
            extension.log("No active workers found. All enabled workers are offline.");
            if (extension.ui?.showToast) {
                extension.ui.showToast(extension.app, "error", "All Workers Offline", 
                    `${enabledWorkers.length} worker(s) enabled but all are offline or unreachable. Check worker connections and try again.`, 5000);
            }
            // Fall back to master-only execution
            return extension.originalQueuePrompt(0, promptWrapper);
        }
        
        extension.log(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`, "debug");
        
        // Find all distributed nodes in the workflow
        const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector");
        const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed");
        const allDistributedNodes = [...collectorNodes, ...upscaleNodes];
        
        // Map original node IDs to truly unique job IDs for this specific run
        const job_id_map = new Map(allDistributedNodes.map(node => [node.id, `${executionPrefix}_${node.id}`]));
        
        // Prepare a separate job queue on the backend for each unique job ID
        const preparePromises = Array.from(job_id_map.values()).map(uniqueId => prepareDistributedJob(extension, uniqueId));
        await Promise.all(preparePromises);

        const jobs = [];
        // Use only active workers
        const participants = ['master', ...activeWorkers.map(w => w.id)];

        for (const participantId of participants) {
            const options = { 
                enabled_worker_ids: activeWorkers.map(w => w.id), 
                workflow: promptWrapper.workflow,
                job_id_map: job_id_map // Pass the map of unique IDs
            };
            
            const jobApiPrompt = prepareApiPromptForParticipant(
                extension, promptWrapper.output, participantId, options
            );
            
            if (participantId === 'master') {
                jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } });
            } else {
                const worker = activeWorkers.find(w => w.id === participantId);
                if (worker) {
                    const job = {
                        type: 'worker',
                        worker,
                        prompt: jobApiPrompt,
                        workflow: promptWrapper.workflow
                    };
                    
                    // Add image references if found for remote workers
                    if (options._imageReferences) {
                        job.imageReferences = options._imageReferences;
                    }
                    
                    jobs.push(job);
                }
            }
        }
        
        const result = await executeJobs(extension, jobs);
        return result;
    } catch (error) {
        extension.log("Parallel execution failed: " + error.message, "error");
        throw error;
    }
}

export function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) {
    let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt));
    const isMaster = participantId === 'master';
    
    // Find all distributed nodes once (before pruning)
    const collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
    const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed");
    const allDistributedNodes = [...collectorNodes, ...upscaleNodes];
    
    // For workers, prune the workflow to only include distributed node dependencies
    if (!isMaster && allDistributedNodes.length > 0) {
        jobApiPrompt = pruneWorkflowForWorker(extension, jobApiPrompt, allDistributedNodes);
    }
    
    // Handle image references for remote workers
    if (!isMaster && options.enabled_worker_ids) {
        // Check if this is a remote worker
        const workerId = participantId;
        const workerInfo = extension.config.workers.find(w => w.id === workerId);
        const isRemote = workerInfo && workerInfo.host;
        
        if (isRemote) {
            // Find all image references in the pruned workflow
            const imageReferences = findImageReferences(extension, jobApiPrompt);
            if (imageReferences.size > 0) {
                extension.log(`Found ${imageReferences.size} image references for remote worker ${workerId}`, "debug");
                // Store image references for later processing
                options._imageReferences = imageReferences;
            }
        }
    }
    
    // Handle Distributed seed nodes
    const distributorNodes = findNodesByClass(jobApiPrompt, "DistributedSeed");
    if (distributorNodes.length > 0) {
        extension.log(`Found ${distributorNodes.length} seed node(s)`, "debug");
    }
    
    for (const seedNode of distributorNodes) {
        const { inputs } = jobApiPrompt[seedNode.id];
        inputs.is_worker = !isMaster;
        if (!isMaster) {
            const workerIndex = options.enabled_worker_ids.indexOf(participantId);
            inputs.worker_id = `worker_${workerIndex}`;
            extension.log(`Set seed node ${seedNode.id} for worker ${workerIndex}`, "debug");
        }
    }
    
    // Handle Distributed collector nodes (already found above)
    for (const collector of collectorNodes) {
        const { inputs } = jobApiPrompt[collector.id];
        
        // Check if this collector is downstream from a distributed upscaler
        const hasUpstreamDistributedUpscaler = hasUpstreamNode(
            jobApiPrompt, 
            collector.id, 
            'UltimateSDUpscaleDistributed'
        );
        
        if (hasUpstreamDistributedUpscaler) {
            // Set pass_through mode for this collector
            inputs.pass_through = true;
            extension.log(`Collector ${collector.id} set to pass-through mode (downstream from distributed upscaler)`, "debug");
        } else {
            // Normal collector behavior
            // Get the unique job ID from the map created for this execution
            const uniqueJobId = options.job_id_map ? options.job_id_map.get(collector.id) : collector.id;
            
            // Use the truly unique ID for this execution
            inputs.multi_job_id = uniqueJobId;
            inputs.is_worker = !isMaster;
            if (isMaster) {
                inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
            } else {
                inputs.master_url = extension.getMasterUrl();
                // Also make the worker_job_id unique to prevent potential caching issues
                inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`;
                inputs.worker_id = participantId;
            }
        }
    }
    
    // Handle Ultimate SD Upscale Distributed nodes
    for (const upscaleNode of upscaleNodes) {
        const { inputs } = jobApiPrompt[upscaleNode.id];
        
        // Get the unique job ID from the map
        const uniqueJobId = options.job_id_map ? options.job_id_map.get(upscaleNode.id) : upscaleNode.id;
        
        inputs.multi_job_id = uniqueJobId;
        inputs.is_worker = !isMaster;
        
        if (isMaster) {
            inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
        } else {
            inputs.master_url = extension.getMasterUrl();
            inputs.worker_id = participantId;
            // Workers also need the enabled_worker_ids to calculate tile distribution
            inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
        }
    }
    
    return jobApiPrompt;
}

export async function prepareDistributedJob(extension, multi_job_id) {
    try {
        await extension.api.prepareJob(multi_job_id);
    } catch (error) {
        extension.log("Error preparing job: " + error.message, "error");
        throw error;
    }
}

export async function executeJobs(extension, jobs) {
    let masterPromptId = null;
    
    // Pre-load all unique images before dispatching to workers
    const allImageReferences = new Map();
    for (const job of jobs) {
        if (job.type === 'worker' && job.imageReferences) {
            for (const [filename, info] of job.imageReferences) {
                allImageReferences.set(filename, info);
            }
        }
    }
    
    if (allImageReferences.size > 0) {
        extension.log(`Pre-loading ${allImageReferences.size} unique image(s) for all workers`, "debug");
        await loadImagesForWorker(extension, allImageReferences);
    }
    
    // Now dispatch jobs in parallel
    const promises = jobs.map(job => {
        if (job.type === 'master') {
            return extension.originalQueuePrompt(0, job.promptWrapper).then(result => {
                masterPromptId = result;
                return result;
            });
        } else {
            return dispatchToWorker(extension, job.worker, job.prompt, job.workflow, job.imageReferences);
        }
    });
    await Promise.all(promises);
    return masterPromptId || { "prompt_id": "distributed-job-dispatched" };
}

async function dispatchToWorker(extension, worker, prompt, workflow, imageReferences) {
    const workerUrl = extension.getWorkerUrl(worker);
    
    // Debug logging - always log to console for debugging
    extension.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`, "debug");
    extension.log('[Distributed] Worker URL: ' + workerUrl, "debug");
    
    // Handle image uploads for remote workers
    if (imageReferences && imageReferences.size > 0) {
        extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug");
        
        try {
            // Load images from master
            const images = await loadImagesForWorker(extension, imageReferences);
            
            // Upload images to worker
            if (images.length > 0) {
                await uploadImagesToWorker(extension, workerUrl, images);
                extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug");
            }
        } catch (error) {
            extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error");
            // Continue with workflow execution even if image upload fails
        }
    }
    
    const promptToSend = {
        prompt,
        extra_data: { extra_pnginfo: { workflow } },
        client_id: api.clientId
    };
    
    extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug");
    
    try {
        await fetch(`${workerUrl}/prompt`, { 
            method: 'POST', 
            headers: { 'Content-Type': 'application/json' }, 
            mode: 'cors',
            body: JSON.stringify(promptToSend) 
        });
    } catch (e) {
        extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error");
    }
}

export async function loadImagesForWorker(extension, imageReferences) {
    const images = [];
    
    // Use a cache to avoid loading the same image multiple times
    if (!extension._imageCache) {
        extension._imageCache = new Map();
    }
    
    for (const [filename, info] of imageReferences) {
        try {
            // Check cache first
            if (extension._imageCache.has(filename)) {
                images.push(extension._imageCache.get(filename));
                extension.log(`Using cached image: ${filename}`, "debug");
                continue;
            }
            
            // Limit cache size
            if (extension._imageCache.size >= 10) {
                const oldestKey = extension._imageCache.keys().next().value;
                extension._imageCache.delete(oldestKey);
                extension.log(`Evicted oldest cache entry: ${oldestKey} (cache limit reached)`, "debug");
            }
            
            // Load image from master's filesystem via API
            try {
                const data = await extension.api.loadImage(filename);
                const imageData = {
                    name: filename,
                    image: data.image_data
                };
                images.push(imageData);
                
                // Cache the image for future use
                extension._imageCache.set(filename, imageData);
                extension.log(`Loaded and cached image: ${filename}`, "debug");
            } catch (loadError) {
                extension.log(`Failed to load image ${filename}: ${loadError.message}`, "error");
                throw loadError;
            }
        } catch (error) {
            extension.log(`Error loading image ${filename}: ${error.message}`, "error");
        }
    }
    
    // Clear cache after a reasonable time to avoid memory issues
    setTimeout(() => {
        if (extension._imageCache && extension._imageCache.size > 0) {
            extension.log(`Clearing image cache (${extension._imageCache.size} images)`, "debug");
            extension._imageCache.clear();
        }
    }, TIMEOUTS.IMAGE_CACHE_CLEAR); // Clear after 30 seconds
    
    return images;
}

export async function uploadImagesToWorker(extension, workerUrl, images) {
    // Upload images to worker's ComfyUI instance
    for (const imageData of images) {
        const formData = new FormData();
        
        // Detect MIME type from base64 header
        const mimeMatch = imageData.image.match(/^data:(image\/\w+);base64,/);
        const mimeType = mimeMatch ? mimeMatch[1] : 'image/png'; // Default to PNG if not detected
        
        // Convert base64 to blob
        const base64Data = imageData.image.replace(/^data:image\/\w+;base64,/, '');
        const byteCharacters = atob(base64Data);
        const byteNumbers = new Array(byteCharacters.length);
        for (let i = 0; i < byteCharacters.length; i++) {
            byteNumbers[i] = byteCharacters.charCodeAt(i);
        }
        const byteArray = new Uint8Array(byteNumbers);
        const blob = new Blob([byteArray], { type: mimeType });
        
        // Use original filename without heavy cleaning
        let cleanName = imageData.name;
        let subfolder = '';
        
        // Extract subfolder if present
        if (cleanName.includes('/')) {
            const parts = cleanName.split('/');
            subfolder = parts.slice(0, -1).join('/');
            cleanName = parts[parts.length - 1];
        }
        
        formData.append('image', blob, cleanName);
        formData.append('type', 'input');
        formData.append('subfolder', subfolder);
        formData.append('overwrite', 'true');
        
        try {
            const response = await fetch(`${workerUrl}/upload/image`, {
                method: 'POST',
                mode: 'cors',
                body: formData
            });
            
            if (!response.ok) {
                throw new Error(`Upload failed: ${response.statusText}`);
            }
            
            extension.log(`Uploaded image to worker: ${imageData.name} -> ${subfolder}/${cleanName}`, "debug");
        } catch (error) {
            extension.log(`Failed to upload ${imageData.name}: ${error.message}`, "error");
            // Continue with other images
        }
    }
}

export async function performPreflightCheck(extension, workers) {
    if (workers.length === 0) return [];
    
    extension.log(`Performing pre-flight health check on ${workers.length} workers...`, "debug");
    const startTime = Date.now();
    
    const checkPromises = workers.map(async (worker) => {
        const url = extension.getWorkerUrl(worker, '/prompt');
        
        extension.log(`Pre-flight checking ${worker.name} at: ${url}`, "debug");
        
        try {
            const response = await fetch(url, {
                method: 'GET',
                mode: 'cors',
                signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
            });
            
            if (response.ok) {
                extension.log(`Worker ${worker.name} is active`, "debug");
                return { worker, active: true };
            } else {
                extension.log(`Worker ${worker.name} returned ${response.status}`, "debug");
                return { worker, active: false };
            }
        } catch (error) {
            extension.log(`Worker ${worker.name} is offline or unreachable: ${error.message}`, "debug");
            return { worker, active: false };
        }
    });
    
    const results = await Promise.all(checkPromises);
    const activeWorkers = results.filter(r => r.active).map(r => r.worker);
    
    const elapsed = Date.now() - startTime;
    extension.log(`Pre-flight check completed in ${elapsed}ms. Active workers: ${activeWorkers.length}/${workers.length}`, "debug");
    
    // Update UI status indicators for inactive workers
    results.filter(r => !r.active).forEach(r => {
        const statusDot = document.getElementById(`status-${r.worker.id}`);
        if (statusDot) {
            // Remove pulsing animation once status is determined
            statusDot.classList.remove('status-pulsing');
            statusDot.style.backgroundColor = "#c04c4c"; // Red for offline
            statusDot.title = "Offline - Cannot connect";
        }
    });
    
    return activeWorkers;
}