code.js

figma.showUI(__html__, { width: 320, height: 400 });

// Flag to track if update is in progress
let shouldCancelUpdate = false;

// When the plugin starts, check if there's a selection
checkSelection();

// Listen for selection changes
figma.on("selectionchange", () => {
  checkSelection();
});

// Function to check the current selection
function checkSelection() {
  const selection = figma.currentPage.selection;
  
  if (selection.length === 0) {
    figma.ui.postMessage({ type: "no-selection" });
  } else {
    figma.ui.postMessage({ 
      type: "selection-ready",
      count: selection.length
    });
  }
}

// Listen for messages from the UI
figma.ui.onmessage = async (msg) => {
  if (msg.type === "update-components") {
    shouldCancelUpdate = false;
    await updateComponentsAndVariablesToLatest();
  } else if (msg.type === "cancel-update") {
    shouldCancelUpdate = true;
    figma.ui.postMessage({ type: "update-cancelled" });
    
    // After 3 seconds, recheck the selection to update the status
    setTimeout(() => {
      checkSelection();
    }, 3000);
  } else if (msg.type === "cancel") {
    figma.closePlugin();
  }
};

// The main function to update components and variables to their latest versions
async function updateComponentsAndVariablesToLatest() {
  const selection = figma.currentPage.selection;
  
  if (selection.length === 0) {
    figma.ui.postMessage({ type: "error", message: "No selection found" });
    return;
  }

  // Track our progress
  let totalComponents = 0;
  let updatedComponents = 0;
  let skippedComponents = 0;
  let errors = 0;
  
  // For variables
  let totalVariables = 0;
  let updatedVariables = 0;
  let skippedVariables = 0;
  let variableErrors = 0;

  // Step 1: Collect all instance nodes
  figma.ui.postMessage({ type: "status", message: "Searching for rogue instances..." });
  const instanceNodes = [];
  const nodesWithVariables = [];
  
  function collectNodesForUpdate(node) {
    try {
      if (!node.visible) return;
      
      // Check if it's a component instance
      if (node.type === "INSTANCE") {
        instanceNodes.push(node);
      }
      
      // Check if the node has variable bindings
      if (node.boundVariables && Object.keys(node.boundVariables).length > 0) {
        nodesWithVariables.push(node);
      }
      
      // Recursively check children
      if ("children" in node) {
        for (const child of node.children) {
          if (!child.locked) {
            collectNodesForUpdate(child);
          }
        }
      }
    } catch (error) {
      console.error("Error collecting nodes:", error);
    }
  }
  
  // Collect all instances and variable-bound nodes from the selection
  try {
    for (const node of selection) {
      if (shouldCancelUpdate) return;
      collectNodesForUpdate(node);
    }
  } catch (error) {
    console.error("Error in selection processing:", error);
    figma.ui.postMessage({ 
      type: "error", 
      message: "Error processing selection: " + error.message 
    });
    return;
  }
  
  totalComponents = instanceNodes.length;
  totalVariables = nodesWithVariables.length;
  
  const totalItems = totalComponents + totalVariables;
  
  if (totalItems === 0) {
    figma.ui.postMessage({ 
      type: "complete", 
      total: 0,
      updated: 0,
      skipped: 0,
      errors: 0
    });
    return;
  }
  
  // Start the update process
  figma.ui.postMessage({ 
    type: "start-update", 
    total: totalItems
  });
  
  // Process components one by one
  for (let i = 0; i < instanceNodes.length; i++) {
    if (shouldCancelUpdate) return;
    
    const node = instanceNodes[i];
    
    try {
      // Check if this is a component that can be updated
      let isRemoteComponent = false;
      let componentKey = null;
      let mainComponent = null;
      
      // Try to get the main component using the async method
      try {
        mainComponent = await node.getMainComponentAsync();
        if (mainComponent && mainComponent.remote) {
          isRemoteComponent = true;
          componentKey = mainComponent.key;
        }
      } catch (error) {
        console.log("Error getting main component:", error);
      }
      
      // If it's a remote component, try to update it
      if (isRemoteComponent && componentKey) {
        try {
          // Try to get the latest version
          const latestVersion = await figma.importComponentByKeyAsync(componentKey);
          
          if (latestVersion && latestVersion !== mainComponent) {
            // Swap to the latest version
            node.swapComponent(latestVersion);
            updatedComponents++;
          } else {
            skippedComponents++;
          }
        } catch (error) {
          console.log("Error importing component:", error);
          skippedComponents++;
        }
      } else {
        skippedComponents++;
      }
    } catch (error) {
      console.error("Error processing component:", error);
      errors++;
    }
    
    // Update progress
    figma.ui.postMessage({ 
      type: "progress", 
      current: i + 1,
      total: totalItems,
      updated: updatedComponents + updatedVariables
    });
    
    // Add a small delay to prevent memory issues
    await new Promise(resolve => setTimeout(resolve, 50));
  }
  
  // Process variables one by one
  for (let i = 0; i < nodesWithVariables.length; i++) {
    if (shouldCancelUpdate) return;
    
    const node = nodesWithVariables[i];
    
    try {
      // Get the variable bindings for this node
      const boundVariables = node.boundVariables;
      
      if (boundVariables) {
        let nodeUpdated = false;
        
        // Check each variable binding
        for (const [property, binding] of Object.entries(boundVariables)) {
          // For direct variable bindings
          if (binding.type === 'VARIABLE') {
            const variable = binding.id;
            
            try {
              // Check if this is a remote variable that can be updated
              if (variable.remote) {
                const variableKey = variable.key;
                
                // Try to get the latest version
                const latestVariable = await figma.variables.importVariableByKeyAsync(variableKey);
                
                if (latestVariable && latestVariable !== variable) {
                  // Update the binding to use the latest variable
                  node.setBoundVariable(property, latestVariable);
                  nodeUpdated = true;
                }
              }
            } catch (error) {
              console.log(`Error updating variable binding for ${property}:`, error);
            }
          }
          // For variable alias bindings
          else if (binding.type === 'VARIABLE_ALIAS') {
            try {
              const aliasVariable = binding.id;
              
              if (aliasVariable.remote) {
                const variableKey = aliasVariable.key;
                
                // Try to get the latest version
                const latestVariable = await figma.variables.importVariableByKeyAsync(variableKey);
                
                if (latestVariable && latestVariable !== aliasVariable) {
                  // Update the binding to use the latest variable
                  node.setBoundVariable(property, {
                    type: 'VARIABLE_ALIAS',
                    id: latestVariable
                  });
                  nodeUpdated = true;
                }
              }
            } catch (error) {
              console.log(`Error updating variable alias binding for ${property}:`, error);
            }
          }
        }
        
        if (nodeUpdated) {
          updatedVariables++;
        } else {
          skippedVariables++;
        }
      } else {
        skippedVariables++;
      }
    } catch (error) {
      console.error("Error processing variable bindings:", error);
      variableErrors++;
    }
    
    // Update progress
    figma.ui.postMessage({ 
      type: "progress", 
      current: instanceNodes.length + i + 1,
      total: totalItems,
      updated: updatedComponents + updatedVariables
    });
    
    // Add a small delay to prevent memory issues
    await new Promise(resolve => setTimeout(resolve, 50));
  }
  
  // Only send complete message if we didn't cancel
  if (!shouldCancelUpdate) {
    // Report the results
    figma.ui.postMessage({ 
      type: "complete", 
      total: totalComponents + totalVariables,
      updated: updatedComponents + updatedVariables,
      skipped: skippedComponents + skippedVariables,
      errors: errors + variableErrors
    });
  }
}