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

/**
 * Convert paths in the API prompt to match the target platform's separator
 * @param {Object} apiPrompt - The workflow API prompt
 * @param {string} targetSeparator - The target path separator ('\\' or '/')
 * @returns {Object} The converted API prompt
 */
function convertPathsForPlatform(apiPrompt, targetSeparator) {
    // Validate target separator
    if (!targetSeparator || (targetSeparator !== '/' && targetSeparator !== '\\')) {
        console.warn('[Distributed] Invalid target separator:', targetSeparator, '- skipping path conversion');
        return apiPrompt;
    }
    
    // Regex to identify likely file paths with extensions
    const isLikelyFilename = (value) => {
        return value.match(/\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)(\s*\[\w+\])?$/i);
    };
    const isImageOrVideo = (value) => {
        return value.match(/\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i);
    };
    
    function convert(obj) {
        if (typeof obj === 'string') {
            // Only convert strings that look like file paths
            if ((obj.includes('\\') || obj.includes('/')) && isLikelyFilename(obj)) {
                const trimmed = obj.trim();
                const hasDrive = /^[A-Za-z]:\\\\|^[A-Za-z]:\//.test(trimmed);
                const isAbsolute = trimmed.startsWith('/') || trimmed.startsWith('\\\\');
                const hasProtocol = /^\w+:\/\//.test(trimmed);

                // For annotated relative image/video paths, keep forward slashes
                if (!hasDrive && !isAbsolute && !hasProtocol && isImageOrVideo(trimmed)) {
                    return trimmed.replace(/[\\\\]/g, '/');
                }
                // Otherwise replace any path separator with the worker's target separator
                return trimmed.replace(/[\\\\\/]/g, targetSeparator);
            }
            return obj;
        } else if (Array.isArray(obj)) {
            return obj.map(convert);
        } else if (typeof obj === 'object' && obj !== null) {
            const newObj = {};
            for (const [key, value] of Object.entries(obj)) {
                newObj[key] = convert(value);
            }
            return newObj;
        }
        return obj;
    }
    
    return convert(apiPrompt);
}

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;
            const hasDistQueue = findNodesByClass(prompt.output, "DistributedQueue").length > 0;

            if (hasCollector || hasDistUpscale) {
                const result = await executeParallelDistributed(extension, prompt);
                // Immediate status check for instant feedback
                extension.checkAllWorkerStatuses();
                // Another check after a short delay to catch state changes
                setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
                return result;
            }

            // DistributedQueue: route entire workflow to least-busy worker
            if (hasDistQueue) {
                const result = await executeQueueDistributed(extension, prompt);
                extension.checkAllWorkerStatuses();
                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");

        // Check if master host might be unreachable by workers (cloudflare tunnel down)
        const masterHost = extension.config?.master?.host || '';
        const isCloudflareHost = /\.(trycloudflare\.com|cloudflare\.dev)$/i.test(masterHost);

        if (isCloudflareHost && activeWorkers.length > 0) {
            // Try to verify if the cloudflare tunnel is actually up
            try {
                const testUrl = `${window.location.protocol}//${masterHost}/prompt`;
                const response = await fetch(testUrl, {
                    method: 'GET',
                    mode: 'cors',
                    cache: 'no-cache',
                    signal: AbortSignal.timeout(3000) // 3 second timeout
                });

                if (!response.ok) {
                    throw new Error('Master not reachable');
                }
            } catch (error) {
                // Cloudflare tunnel appears to be down
                extension.log(`Master host ${masterHost} is not reachable - cloudflare tunnel may be down`, "error");

                if (extension.ui?.showCloudflareWarning) {
                    extension.ui.showCloudflareWarning(extension, masterHost);
                }

                // Stop execution - workers won't be able to send results back
                extension.log("Blocking execution - workers cannot reach master at cloudflare domain", "error");
                return null; // This will prevent the workflow from running
            }
        }

        // Find all distributed nodes in the workflow
        const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector");
        const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed");
        const allDistributedNodes = [...collectorNodes, ...upscaleNodes];

        const wantsDelegate = Boolean(extension.config?.settings?.master_delegate_only);
        let masterDelegateActive = false;
        if (wantsDelegate) {
            if (activeWorkers.length === 0) {
                extension.log("Master delegate-only mode disabled: no active workers detected. Falling back to master execution.", "debug");
            } else if (upscaleNodes.length > 0) {
                extension.log("Master delegate-only mode is not yet supported for UltimateSDUpscaleDistributed nodes. Falling back to normal execution.", "warn");
            } else if (collectorNodes.length === 0) {
                extension.log("Master delegate-only mode requested but no DistributedCollector nodes found. Falling back to normal execution.", "debug");
            } else {
                masterDelegateActive = true;
                extension.log("Master delegate-only mode enabled: master will orchestrate without running upstream nodes.", "debug");
            }
        }
        
        // 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
                delegate_master: masterDelegateActive && participantId === 'master'
            };
            
            const jobApiPrompt = await 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;
    }
}

/**
 * Execute a workflow via DistributedQueue - routes to the least-busy worker.
 * The workflow runs ONLY on the selected worker, not on the master.
 */
export async function executeQueueDistributed(extension, promptWrapper) {
    try {
        const enabledWorkers = extension.enabledWorkers;

        if (enabledWorkers.length === 0) {
            extension.log("DistributedQueue: No enabled workers. Falling back to local execution.", "warn");
            return extension.originalQueuePrompt(0, promptWrapper);
        }

        // Fetch queue status from all enabled workers
        const workerStatuses = await fetchWorkerQueueStatuses(extension, enabledWorkers);

        if (workerStatuses.length === 0) {
            extension.log("DistributedQueue: No reachable workers. Falling back to local execution.", "warn");
            return extension.originalQueuePrompt(0, promptWrapper);
        }

        // Select the least-busy worker
        const selectedWorker = selectLeastBusyWorker(extension, workerStatuses);
        const worker = selectedWorker.worker;
        const queueRemaining = selectedWorker.queueRemaining;

        extension.log(`DistributedQueue: Routing to ${worker.name} (queue_remaining=${queueRemaining})`, "info");

        // Prepare the prompt with skip_dispatch=True for DistributedQueue nodes
        const promptToSend = JSON.parse(JSON.stringify(promptWrapper.output));
        markSkipDispatch(promptToSend);

        // Dispatch to the selected worker
        const workerUrl = extension.getWorkerUrl(worker);
        try {
            const response = await fetch(`${workerUrl}/prompt`, {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                mode: 'cors',
                body: JSON.stringify({
                    prompt: promptToSend,
                    extra_data: { extra_pnginfo: { workflow: promptWrapper.workflow } },
                    client_id: api.clientId
                })
            });

            if (!response.ok) {
                throw new Error(`Worker returned ${response.status}`);
            }

            const result = await response.json();
            extension.log(`DistributedQueue: Dispatched to ${worker.name}, prompt_id=${result.prompt_id}`, "info");

            return { prompt_id: result.prompt_id, worker_id: worker.id };
        } catch (error) {
            extension.log(`DistributedQueue: Failed to dispatch to ${worker.name}: ${error.message}`, "error");
            // Fall back to local execution
            return extension.originalQueuePrompt(0, promptWrapper);
        }
    } catch (error) {
        extension.log("DistributedQueue execution failed: " + error.message, "error");
        throw error;
    }
}

async function fetchWorkerQueueStatuses(extension, workers) {
    const statuses = [];

    const checkPromises = workers.map(async (worker) => {
        const url = extension.getWorkerUrl(worker, '/prompt');
        try {
            const response = await fetch(url, {
                method: 'GET',
                mode: 'cors',
                cache: 'no-store',
                signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
            });

            if (response.ok) {
                const data = await response.json();
                const queueRemaining = data.exec_info?.queue_remaining || 0;
                return { worker, queueRemaining, online: true };
            }
        } catch (error) {
            extension.log(`DistributedQueue: Worker ${worker.name} unreachable: ${error.message}`, "debug");
        }
        return null;
    });

    const results = await Promise.all(checkPromises);
    return results.filter(r => r !== null);
}

function selectLeastBusyWorker(extension, statuses) {
    // Find idle workers (queue_remaining == 0)
    const idleWorkers = statuses.filter(s => s.queueRemaining === 0);

    if (idleWorkers.length > 0) {
        // Round-robin among idle workers
        if (!extension._distributedQueueRRIndex) {
            extension._distributedQueueRRIndex = 0;
        }
        const index = extension._distributedQueueRRIndex % idleWorkers.length;
        extension._distributedQueueRRIndex++;
        return idleWorkers[index];
    }

    // No idle workers - pick the one with the shortest queue
    return statuses.reduce((min, s) => s.queueRemaining < min.queueRemaining ? s : min);
}

function markSkipDispatch(promptObj) {
    for (const node of Object.values(promptObj)) {
        if (node && typeof node === 'object' && node.class_type === 'DistributedQueue') {
            node.inputs = node.inputs || {};
            node.inputs.skip_dispatch = true;
        }
    }
}

export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) {
    let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt));
    const isMaster = participantId === 'master';
    const delegateMaster = Boolean(options.delegate_master);
    
    // Find all distributed nodes once (before pruning)
    let collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
    const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed");
    let allDistributedNodes = [...collectorNodes, ...upscaleNodes];

    if (isMaster && delegateMaster) {
        if (upscaleNodes.length > 0) {
            extension.log("Delegate-only master mode does not support UltimateSDUpscaleDistributed nodes yet. Using full prompt.", "warn");
        } else if (collectorNodes.length === 0) {
            extension.log("Delegate-only master mode requested but no collectors found in master prompt. Using full prompt.", "debug");
        } else {
            jobApiPrompt = prepareMasterDelegatePrompt(extension, jobApiPrompt, collectorNodes);
            collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
            allDistributedNodes = [...collectorNodes, ...upscaleNodes];
        }
    }
    
    // For workers, handle platform-specific path conversion
    if (!isMaster) {
        const workerInfo = extension.config.workers.find(w => w.id === participantId);
        
        if (workerInfo && workerInfo.host) {
            // Remote or cloud worker - needs path translation
            try {
                const workerUrl = extension.getWorkerUrl(workerInfo);
                const systemInfo = await getCachedWorkerSystemInfo(workerUrl);
                const targetSeparator = systemInfo?.platform?.path_separator;
                
                if (targetSeparator) {
                    // Convert paths to match worker's platform
                    jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator);
                    extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug");
                } else {
                    extension.log(`No path separator found for worker ${participantId}, skipping path conversion`, "debug");
                }
            } catch (e) {
                extension.log(`Failed to get system info for worker ${participantId}: ${e.message}`, "warn");
                // Continue without path conversion
            }
        }
        
        // Prune the workflow to only include distributed node dependencies
        if (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/video references in the pruned workflow
            const imageReferences = findImageReferences(extension, jobApiPrompt);
            if (imageReferences.size > 0) {
                extension.log(`Found ${imageReferences.size} media references (images/videos) 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 || []);
                if (delegateMaster) {
                    inputs.delegate_only = true;
                }
            } 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;
    }
}

function createNumericIdGenerator(promptObj) {
    let maxId = 0;
    for (const key of Object.keys(promptObj)) {
        const numeric = parseInt(key, 10);
        if (!Number.isNaN(numeric)) {
            maxId = Math.max(maxId, numeric);
        }
    }
    return () => {
        maxId += 1;
        return String(maxId);
    };
}

function prepareMasterDelegatePrompt(extension, apiPrompt, collectorNodes) {
    const collectorIds = collectorNodes.map((node) => node.id);
    const nodesToKeep = findCollectorDownstreamNodes(apiPrompt, collectorIds);

    // Ensure collectors themselves are present
    collectorIds.forEach((id) => nodesToKeep.add(id));

    const prunedPrompt = {};
    nodesToKeep.forEach((nodeId) => {
        const nodeData = apiPrompt[nodeId];
        if (nodeData) {
            prunedPrompt[nodeId] = JSON.parse(JSON.stringify(nodeData));
        }
    });

    const prunedIds = new Set(Object.keys(prunedPrompt));

    // Remove dangling references to trimmed nodes
    for (const [nodeId, node] of Object.entries(prunedPrompt)) {
        if (!node.inputs) continue;
        for (const [inputName, inputValue] of Object.entries(node.inputs)) {
            if (Array.isArray(inputValue) && inputValue.length === 2) {
                const sourceId = String(inputValue[0]);
                if (!prunedIds.has(sourceId)) {
                    delete node.inputs[inputName];
                    extension.log(`Removed upstream reference '${inputName}' from node ${nodeId} for delegate-only master prompt.`, "debug");
                }
            }
        }
    }

    const nextId = createNumericIdGenerator(prunedPrompt);
    collectorIds.forEach((collectorId) => {
        const collectorEntry = prunedPrompt[collectorId];
        if (!collectorEntry) return;

        const placeholderId = nextId();
        prunedPrompt[placeholderId] = {
            class_type: "DistributedEmptyImage",
            inputs: {
                height: 64,
                width: 64,
                channels: 3
            }
        };
        prunedIds.add(placeholderId);

        collectorEntry.inputs = collectorEntry.inputs || {};
        collectorEntry.inputs.images = [placeholderId, 0];
    });

    extension.log(`Prepared delegate-only master prompt with ${Object.keys(prunedPrompt).length} nodes (kept ${nodesToKeep.size}).`, "debug");
    return prunedPrompt;
}

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 media file(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);
    
    // Trigger immediate status check for instant feedback
    extension.checkAllWorkerStatuses();
    
    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) {
        // Check if this is a local worker (same host as master)
        const isLocalWorker = workerUrl.includes('127.0.0.1') || workerUrl.includes('localhost');
        
        if (isLocalWorker) {
            extension.log(`[Distributed] Skipping image processing for local worker ${worker.name} (shares filesystem with master)`, "debug");
        } else {
            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,
        workflow,
        client_id: api.clientId
    };

    extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug");

    try {
        if (extension.config?.settings?.websocket_orchestration) {
            const wsUrl = buildWorkerWebSocketUrl(workerUrl);
            await dispatchPromptViaWebSocket(extension, wsUrl, promptToSend);
            return;
        }
        await fetch(`${workerUrl}/prompt`, {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            mode: 'cors',
            body: JSON.stringify({
                prompt,
                extra_data: { extra_pnginfo: { workflow } },
                client_id: api.clientId
            })
        });
    } catch (e) {
        extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error");
    }
}

function buildWorkerWebSocketUrl(workerUrl) {
    const url = new URL(workerUrl);
    url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:';
    url.pathname = '/distributed/worker_ws';
    url.search = '';
    return url.toString();
}

function generateRequestId() {
    if (crypto.randomUUID) {
        return crypto.randomUUID();
    }
    return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => {
        const r = Math.random() * 16 | 0;
        const v = c === 'x' ? r : (r & 0x3 | 0x8);
        return v.toString(16);
    });
}

async function dispatchPromptViaWebSocket(extension, wsUrl, payload) {
    const requestId = generateRequestId();
    const message = {
        type: 'dispatch_prompt',
        request_id: requestId,
        prompt: payload.prompt,
        workflow: payload.workflow,
        client_id: payload.client_id
    };

    await new Promise((resolve, reject) => {
        const ws = new WebSocket(wsUrl);
        const timeoutId = setTimeout(() => {
            ws.close();
            reject(new Error('Websocket dispatch timed out.'));
        }, 60000);

        ws.addEventListener('open', () => {
            ws.send(JSON.stringify(message));
        });

        ws.addEventListener('message', (event) => {
            try {
                const data = JSON.parse(event.data || '{}');
                if (data.type !== 'dispatch_ack' || data.request_id !== requestId) {
                    return;
                }
                clearTimeout(timeoutId);
                ws.close();
                if (data.ok) {
                    resolve();
                } else {
                    reject(new Error(data.error || 'Worker rejected websocket dispatch.'));
                }
            } catch (error) {
                clearTimeout(timeoutId);
                ws.close();
                reject(error);
            }
        });

        ws.addEventListener('error', () => {
            clearTimeout(timeoutId);
            ws.close();
            reject(new Error('Websocket connection failed.'));
        });
    });
    extension.log(`[Distributed] Websocket dispatch succeeded for ${wsUrl}`, "debug");
}

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,
                    hash: data.hash  // Include hash from the response
                };
                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) {
        // Check if file already exists with matching hash
        if (imageData.hash) {
            try {
                const checkResponse = await fetch(`${workerUrl}/distributed/check_file`, {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    mode: 'cors',
                    body: JSON.stringify({ 
                        filename: imageData.name, 
                        hash: imageData.hash 
                    })
                });
                
                if (checkResponse.ok) {
                    const result = await checkResponse.json();
                    if (result.exists && result.hash_matches) {
                        extension.log(`File ${imageData.name} already exists on worker with matching hash, skipping upload`, "debug");
                        continue;
                    }
                }
            } catch (error) {
                // If check fails, proceed with upload
                extension.log(`Failed to check file existence for ${imageData.name}: ${error.message}`, "debug");
            }
        }
        
        const formData = new FormData();
        
        // Detect MIME type from base64 header (supports both image and video)
        const mimeMatch = imageData.image.match(/^data:((?:image|video)\/\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|video)\/\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 (handle both slash styles)
        if (cleanName.includes('/') || cleanName.includes('\\')) {
            const parts = cleanName.replace(/\\/g, '/').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',
                cache: 'no-store',
                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) {
            if (error?.name === 'AbortError') {
                extension.log(`Worker ${worker.name} pre-flight check timed out; assuming active`, "debug");
                return { worker, active: true, uncertain: true };
            }
            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;
}