import { app } from "../../scripts/app.js";


const NODE_CLASS = "LoraRangeLoaderModelOnly";


function getWidget(node, name) {
  return node.widgets?.find((widget) => widget.name === name);
}


function getSortedLoras(currentWidget) {
  const values = Array.isArray(currentWidget?.options?.values) ? [...currentWidget.options.values] : [];
  values.sort((left, right) => String(left).localeCompare(String(right), undefined, { numeric: true, sensitivity: "base" }));
  return values;
}


function getAllowedRange(node) {
  const firstWidget = getWidget(node, "first");
  const lastWidget = getWidget(node, "last");
  const currentWidget = getWidget(node, "current");
  if (!firstWidget || !lastWidget || !currentWidget) {
    return null;
  }

  const sorted = getSortedLoras(currentWidget);
  const firstIndex = sorted.indexOf(firstWidget.value);
  const lastIndex = sorted.indexOf(lastWidget.value);
  if (firstIndex === -1 || lastIndex === -1 || firstIndex > lastIndex) {
    return null;
  }

  return sorted.slice(firstIndex, lastIndex + 1);
}


function advanceCurrent(node) {
  const currentWidget = getWidget(node, "current");
  const modeWidget = getWidget(node, "mode");
  if (!currentWidget || !modeWidget) {
    return;
  }

  const allowedRange = getAllowedRange(node);
  if (!allowedRange || allowedRange.length === 0) {
    return;
  }

  const currentIndex = allowedRange.indexOf(currentWidget.value);
  if (currentIndex === -1) {
    return;
  }

  let nextValue = currentWidget.value;
  switch (modeWidget.value) {
    case "fixed":
      return;
    case "increment":
      nextValue = allowedRange[(currentIndex + 1) % allowedRange.length];
      break;
    case "decrement":
      nextValue = allowedRange[(currentIndex - 1 + allowedRange.length) % allowedRange.length];
      break;
    case "random":
      nextValue = allowedRange[Math.floor(Math.random() * allowedRange.length)];
      break;
    default:
      return;
  }

  currentWidget.value = nextValue;
  currentWidget.callback?.(nextValue);
  node.setDirtyCanvas(true, true);
}


function installAfterQueuedHandler(node) {
  if (node.__loraRangeAfterQueuedInstalled) {
    return;
  }

  const currentWidget = getWidget(node, "current");
  if (!currentWidget) {
    return;
  }

  const originalAfterQueued = currentWidget.afterQueued?.bind(currentWidget);
  currentWidget.afterQueued = function afterQueued(event = {}) {
    originalAfterQueued?.(event);
    if (!event.isPartialExecution) {
      advanceCurrent(node);
    }
  };

  node.__loraRangeAfterQueuedInstalled = true;
}


app.registerExtension({
  name: "Comfy.LoraRangeLoader",
  nodeCreated(node) {
    if (node.constructor.comfyClass !== NODE_CLASS) {
      return;
    }

    installAfterQueuedHandler(node);
  },
});