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

/**
 * Workflow Iterator - ComfyUI JavaScript Extension
 *
 * Provides:
 * 1. Node target browser for WI Parameter nodes
 * 2. Progress bar overlay for WorkflowIterator node
 * 3. Value preview for WI Parameter nodes
 */

const WI_PARAM_TYPES = [
  "WIParameterInt",
  "WIParameterFloat",
  "WIParameterString",
  "WIParameterCombo",
  "WIParameterSeed",
];

app.registerExtension({
  name: "WorkflowIterator",

  async beforeRegisterNodeDef(nodeType, nodeData, app) {
    if (!nodeData.name) return;

    // WorkflowIterator main node - add progress overlay
    if (nodeData.name === "WorkflowIterator") {
      const origOnDrawForeground = nodeType.prototype.onDrawForeground;

      nodeType.prototype.onDrawForeground = function (ctx) {
        if (origOnDrawForeground) {
          origOnDrawForeground.apply(this, arguments);
        }
        drawProgressOverlay(this, ctx);
      };

      const origOnExecuted = nodeType.prototype.onExecuted;
      nodeType.prototype.onExecuted = function (output) {
        if (origOnExecuted) {
          origOnExecuted.apply(this, arguments);
        }
        // Start polling for batch status
        startBatchPolling(this);
      };
    }

    // WI Parameter nodes - add browse button and value preview
    if (WI_PARAM_TYPES.includes(nodeData.name)) {
      const origOnNodeCreated = nodeType.prototype.onNodeCreated;

      nodeType.prototype.onNodeCreated = function () {
        if (origOnNodeCreated) {
          origOnNodeCreated.apply(this, arguments);
        }
        addBrowseButton(this);
        addValuePreview(this);
      };
    }
  },
});

// ============================================================
// Progress Overlay for WorkflowIterator
// ============================================================

const batchStatusCache = new Map();

function drawProgressOverlay(node, ctx) {
  const status = batchStatusCache.get(node.id);
  if (!status || status.total <= 1) return;

  const progress = status.completed / status.total;
  const barHeight = 6;
  const barY = node.size[1] - barHeight - 4;
  const barWidth = node.size[0] - 20;
  const barX = 10;

  // Background
  ctx.fillStyle = "rgba(0, 0, 0, 0.5)";
  ctx.fillRect(barX, barY, barWidth, barHeight);

  // Progress fill
  ctx.fillStyle = status.cancelled
    ? "#ff4444"
    : progress >= 1
    ? "#44ff44"
    : "#4488ff";
  ctx.fillRect(barX, barY, barWidth * progress, barHeight);

  // Border
  ctx.strokeStyle = "rgba(255, 255, 255, 0.3)";
  ctx.strokeRect(barX, barY, barWidth, barHeight);

  // Text
  const text = status.cancelled
    ? `Cancelled (${status.completed}/${status.total})`
    : `${status.completed}/${status.total}`;
  ctx.fillStyle = "#ffffff";
  ctx.font = "11px sans-serif";
  ctx.textAlign = "center";
  ctx.fillText(text, node.size[0] / 2, barY - 4);
  ctx.textAlign = "left";
}

let pollInterval = null;

function startBatchPolling(node) {
  if (pollInterval) return;

  pollInterval = setInterval(async () => {
    try {
      const resp = await api.fetchApi("/wi/active_batches");
      if (resp.ok) {
        const batches = await resp.json();
        const entries = Object.entries(batches);

        if (entries.length === 0) {
          clearInterval(pollInterval);
          pollInterval = null;
          batchStatusCache.delete(node.id);
          app.graph.setDirtyCanvas(true);
          return;
        }

        // Use the most recent active batch
        const [batchId, status] = entries[entries.length - 1];
        batchStatusCache.set(node.id, status);
        app.graph.setDirtyCanvas(true);

        // Stop polling if complete or cancelled
        if (status.completed >= status.total || status.cancelled) {
          setTimeout(() => {
            clearInterval(pollInterval);
            pollInterval = null;
          }, 3000);
        }
      }
    } catch (e) {
      // Silently ignore polling errors
    }
  }, 1000);
}

// ============================================================
// Browse Button for WI Parameter Nodes
// ============================================================

function addBrowseButton(node) {
  const browseWidget = node.addWidget(
    "button",
    "Browse Target Node",
    null,
    () => {
      showNodeBrowser(node);
    }
  );
  browseWidget.serialize = false;
}

function showNodeBrowser(paramNode) {
  // Get all nodes in the graph
  const allNodes = app.graph._nodes;
  if (!allNodes || allNodes.length === 0) return;

  // Build list of unique node titles/types
  const nodeOptions = [];
  const seen = new Set();

  for (const node of allNodes) {
    if (node === paramNode) continue;
    // Skip other WI nodes
    if (
      node.type &&
      (node.type.startsWith("WIParameter") ||
        node.type === "WorkflowIterator")
    )
      continue;

    const title = node.title || node.type || `Node #${node.id}`;
    const key = `${title}__${node.id}`;
    if (seen.has(key)) continue;
    seen.add(key);

    nodeOptions.push({
      title: title,
      type: node.type || "",
      id: node.id,
      node: node,
    });
  }

  // Sort by title
  nodeOptions.sort((a, b) => a.title.localeCompare(b.title));

  // Create modal dialog
  const dialog = document.createElement("dialog");
  dialog.style.cssText = `
    background: #1a1a2e; color: #eee; border: 1px solid #444;
    border-radius: 8px; padding: 16px; min-width: 350px;
    max-height: 70vh; font-family: sans-serif;
  `;

  const title = document.createElement("h3");
  title.textContent = "Select Target Node";
  title.style.cssText = "margin: 0 0 12px 0; font-size: 14px;";
  dialog.appendChild(title);

  // Search input
  const searchInput = document.createElement("input");
  searchInput.type = "text";
  searchInput.placeholder = "Filter nodes...";
  searchInput.style.cssText = `
    width: 100%; padding: 6px; margin-bottom: 8px;
    background: #2a2a4e; color: #eee; border: 1px solid #555;
    border-radius: 4px; box-sizing: border-box;
  `;
  dialog.appendChild(searchInput);

  // Node list
  const listContainer = document.createElement("div");
  listContainer.style.cssText = `
    max-height: 300px; overflow-y: auto; margin-bottom: 12px;
  `;

  function renderList(filter) {
    listContainer.innerHTML = "";
    const filtered = filter
      ? nodeOptions.filter(
          (n) =>
            n.title.toLowerCase().includes(filter.toLowerCase()) ||
            n.type.toLowerCase().includes(filter.toLowerCase())
        )
      : nodeOptions;

    for (const opt of filtered) {
      const item = document.createElement("div");
      item.style.cssText = `
        padding: 6px 8px; cursor: pointer; border-radius: 4px;
        margin: 2px 0; font-size: 13px;
      `;
      item.innerHTML = `<strong>${opt.title}</strong> <span style="color:#888">(${opt.type} #${opt.id})</span>`;

      item.onmouseenter = () => (item.style.background = "#333");
      item.onmouseleave = () => (item.style.background = "");
      item.onclick = () => {
        dialog.close();
        selectNodeTarget(paramNode, opt);
      };

      listContainer.appendChild(item);
    }
  }

  renderList("");
  searchInput.addEventListener("input", (e) => renderList(e.target.value));

  dialog.appendChild(listContainer);

  // Cancel button
  const cancelBtn = document.createElement("button");
  cancelBtn.textContent = "Cancel";
  cancelBtn.style.cssText = `
    padding: 6px 16px; background: #444; color: #eee;
    border: none; border-radius: 4px; cursor: pointer;
  `;
  cancelBtn.onclick = () => dialog.close();
  dialog.appendChild(cancelBtn);

  document.body.appendChild(dialog);
  dialog.showModal();
  searchInput.focus();

  dialog.addEventListener("close", () => dialog.remove());
}

function selectNodeTarget(paramNode, selectedNode) {
  // Set the target_node_title widget
  const titleWidget = paramNode.widgets?.find(
    (w) => w.name === "target_node_title"
  );
  if (titleWidget) {
    titleWidget.value = selectedNode.title;
  }

  // Now show widget picker for the selected node
  showWidgetPicker(paramNode, selectedNode.node);
}

function showWidgetPicker(paramNode, targetNode) {
  if (!targetNode.widgets || targetNode.widgets.length === 0) return;

  const dialog = document.createElement("dialog");
  dialog.style.cssText = `
    background: #1a1a2e; color: #eee; border: 1px solid #444;
    border-radius: 8px; padding: 16px; min-width: 300px;
    font-family: sans-serif;
  `;

  const title = document.createElement("h3");
  title.textContent = `Select Widget on "${targetNode.title || targetNode.type}"`;
  title.style.cssText = "margin: 0 0 12px 0; font-size: 14px;";
  dialog.appendChild(title);

  const listContainer = document.createElement("div");
  listContainer.style.cssText = "max-height: 300px; overflow-y: auto;";

  for (const widget of targetNode.widgets) {
    if (!widget.name || widget.name.startsWith("_")) continue;

    const item = document.createElement("div");
    item.style.cssText = `
      padding: 6px 8px; cursor: pointer; border-radius: 4px;
      margin: 2px 0; font-size: 13px;
    `;

    const widgetType = widget.type || "unknown";
    const currentValue =
      widget.value !== undefined ? ` = ${widget.value}` : "";
    item.innerHTML = `<strong>${widget.name}</strong> <span style="color:#888">(${widgetType}${currentValue})</span>`;

    item.onmouseenter = () => (item.style.background = "#333");
    item.onmouseleave = () => (item.style.background = "");
    item.onclick = () => {
      const widgetNameWidget = paramNode.widgets?.find(
        (w) => w.name === "target_widget_name"
      );
      if (widgetNameWidget) {
        widgetNameWidget.value = widget.name;
      }
      dialog.close();
    };

    listContainer.appendChild(item);
  }

  dialog.appendChild(listContainer);

  const cancelBtn = document.createElement("button");
  cancelBtn.textContent = "Cancel";
  cancelBtn.style.cssText = `
    padding: 6px 16px; background: #444; color: #eee;
    border: none; border-radius: 4px; cursor: pointer; margin-top: 8px;
  `;
  cancelBtn.onclick = () => dialog.close();
  dialog.appendChild(cancelBtn);

  document.body.appendChild(dialog);
  dialog.showModal();
  dialog.addEventListener("close", () => dialog.remove());
}

// ============================================================
// Value Preview for WI Parameter Nodes
// ============================================================

function addValuePreview(node) {
  const previewWidget = node.addWidget("text", "preview", "", () => {}, {
    serialize: false,
    multiline: false,
  });
  previewWidget.computeSize = () => [0, 20];

  // Override draw to show parsed value count
  const origDraw = node.onDrawForeground;
  node.onDrawForeground = function (ctx) {
    if (origDraw) origDraw.apply(this, arguments);

    const valuesWidget = this.widgets?.find((w) => w.name === "values");
    if (!valuesWidget) return;

    const raw = valuesWidget.value || "";
    const typeWidget = this.type;
    let count = 0;

    try {
      if (raw.includes("-") && raw.includes(",") && !raw.includes("\n")) {
        // Range syntax - estimate count
        const parts = raw.split(",").map((s) => s.trim());
        if (parts.length === 2) {
          const rangeParts = parts[0].split("-");
          if (rangeParts.length === 2) {
            const start = parseFloat(rangeParts[0]);
            const end = parseFloat(rangeParts[1]);
            const step = parseFloat(parts[1]);
            if (!isNaN(start) && !isNaN(end) && !isNaN(step) && step > 0) {
              count = Math.floor(Math.abs(end - start) / step) + 1;
            }
          }
        }
      } else if (raw.includes("\n")) {
        count = raw
          .split("\n")
          .filter((l) => l.trim()).length;
      } else {
        count = raw
          .split(",")
          .filter((p) => p.trim()).length;
      }
    } catch (e) {
      count = 0;
    }

    if (count > 0) {
      const text = `${count} value${count !== 1 ? "s" : ""}`;
      const y = this.size[1] - 18;
      ctx.fillStyle = "rgba(100, 180, 255, 0.7)";
      ctx.font = "10px sans-serif";
      ctx.textAlign = "right";
      ctx.fillText(text, this.size[0] - 10, y);
      ctx.textAlign = "left";
    }
  };
}