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

// Image Chooser Gate (Batch) — pauses a prompt, displays every image in the
// incoming batch, and resumes with the selected subset as a new IMAGE batch.

const NODE = "ImageChooserGate";
const R = "/datasete_image_chooser";

const MIN_W = 420;
const MIN_GRID_H = 190;
const TOOLBAR_H = 72;
const MARGIN = 10;
const chooserNodes = new Set();

function graphNodeById(graph, id) {
  if (!graph?.getNodeById) return null;
  return graph.getNodeById(id) ?? graph.getNodeById(parseInt(id, 10));
}

function chooserByLocator(locator) {
  const direct = graphNodeById(app.graph, locator);
  if (direct?.type === NODE) return direct;

  // Native subgraphs identify executing nodes as colon-separated locators:
  // root subgraph-node id -> nested node id -> ... -> chooser id.
  let graph = app.rootGraph ?? app.graph?.rootGraph ?? app.graph;
  let node = null;
  for (const id of String(locator).split(":")) {
    node = graphNodeById(graph, id);
    if (!node) return null;
    graph = node.subgraph;
  }
  return node?.type === NODE ? node : null;
}

function widgetFloor(node) {
  return node._icgState === "idle" ? 0 : 2 * MARGIN + MIN_GRID_H + TOOLBAR_H;
}

function syncWidgetWidth(node) {
  if (node._icgWidget) node._icgWidget.width = node.size?.[0] || MIN_W;
}

function resizeChooser(node) {
  const shown = node._icgState !== "idle";
  if (node._icg?.wrap) node._icg.wrap.style.display = shown ? "flex" : "none";
  const width = node.size?.[0] || MIN_W;
  const target = shown
    ? Math.max(node.size?.[1] || 0, node.computeSize()[1])
    : node.computeSize()[1];
  node.setSize([width, target]);
  syncWidgetWidth(node);
  node.setDirtyCanvas?.(true, true);
}

async function postChooser(node, fields) {
  const form = new FormData();
  form.append("id", String(node._icgWaitId ?? node.id));
  form.append("token", String(node._icgToken ?? ""));
  for (const [key, value] of Object.entries(fields)) form.append(key, value);
  const response = await api.fetchApi(`${R}/select`, { method: "POST", body: form });
  if (!response.ok) throw new Error(`chooser request failed (${response.status})`);
}

async function queueFromHere() {
  const command = app.extensionManager?.command;
  if (command?.execute) {
    try {
      await command.execute("Comfy.QueuePrompt");
      return true;
    } catch (e) { /* use the legacy path below */ }
  }
  try {
    await app.queuePrompt(0, 1);
    return true;
  } catch (e) {
    try {
      await app.queuePrompt(0);
      return true;
    } catch (fallbackError) {
      console.error("[image-chooser] queue failed", fallbackError);
      return false;
    }
  }
}

function selectedIndices(node) {
  return [...node._icgSelected].sort((a, b) => a - b);
}

function updateCells(node) {
  for (const [index, cell] of node._icgCells) {
    const selected = node._icgSelected.has(index);
    cell.classList.toggle("icg-selected", selected);
    cell.setAttribute("aria-pressed", selected ? "true" : "false");
    const check = cell.querySelector(".icg-check");
    if (check) check.textContent = selected ? "✓" : "";
  }
}

function select(node, indices) {
  if (node._icgState !== "paused" || node._icgBusy) return;
  node._icgSelected = new Set(indices);
  node._icgError = "";
  updateCells(node);
  renderToolbar(node);
}

function button(text, className, onclick) {
  const element = document.createElement("button");
  element.type = "button";
  element.textContent = text;
  element.className = className;
  element.onclick = onclick;
  return element;
}

function renderToolbar(node) {
  const { toolbar } = node._icg;
  toolbar.innerHTML = "";
  const count = node._icgItems.length;
  const chosen = selectedIndices(node);

  if (node._icgState === "paused") {
    const all = button("Select all", "icg-secondary", () => {
      select(node, node._icgItems.map((item) => item.index));
    });
    all.disabled = node._icgBusy || chosen.length === count;

    const clear = button("Clear", "icg-secondary", () => select(node, []));
    clear.disabled = node._icgBusy || chosen.length === 0;

    const pass = button(
      chosen.length ? `▶ Pass selected (${chosen.length})` : "▶ Pass selected",
      "icg-pass",
      async () => {
        if (!chosen.length || node._icgBusy) return;
        const revision = node._icgRevision;
        node._icgBusy = true;
        node._icgError = "";
        renderGridState(node);
        renderToolbar(node);
        try {
          await postChooser(node, { selection: JSON.stringify(chosen) });
          if (node._icgRevision === revision) {
            node._icgState = "resolved";
            node._icgResolvedCount = chosen.length;
          }
        } catch (error) {
          console.error("[image-chooser] pass failed", error);
          if (node._icgRevision === revision) {
            node._icgError = "Could not pass the selection";
          }
        } finally {
          if (node._icgRevision !== revision) return;
          node._icgBusy = false;
          renderGridState(node);
          renderToolbar(node);
          resizeChooser(node);
        }
      },
    );
    pass.disabled = node._icgBusy || chosen.length === 0;

    const stop = button("■ Stop", "icg-stop", async () => {
      if (node._icgBusy) return;
      const revision = node._icgRevision;
      node._icgBusy = true;
      node._icgError = "";
      renderGridState(node);
      renderToolbar(node);
      try {
        await postChooser(node, { action: "cancel" });
        if (node._icgRevision === revision) {
          node._icgState = "resolved";
          node._icgStopped = true;
        }
      } catch (error) {
        console.error("[image-chooser] stop failed", error);
        if (node._icgRevision === revision) {
          node._icgError = "Could not stop the run";
        }
      } finally {
        if (node._icgRevision !== revision) return;
        node._icgBusy = false;
        renderGridState(node);
        renderToolbar(node);
      }
    });
    stop.disabled = node._icgBusy;

    const status = document.createElement("span");
    status.className = `icg-status${node._icgError ? " icg-error" : ""}`;
    status.textContent = node._icgError || (node._icgBusy
      ? "sending…"
      : `${chosen.length} of ${count} selected`);

    toolbar.appendChild(all);
    toolbar.appendChild(clear);
    toolbar.appendChild(pass);
    toolbar.appendChild(status);
    toolbar.appendChild(stop);
  } else if (node._icgState === "resolved") {
    const status = document.createElement("span");
    status.className = `icg-status icg-resolved${node._icgError ? " icg-error" : ""}`;
    status.textContent = node._icgError || (node._icgQueueBusy
      ? "queuing…"
      : node._icgStopped
        ? "■ run stopped"
        : `✓ passed ${node._icgResolvedCount} of ${count}`);
    const run = button("▶ Run from here", "icg-run", async () => {
      if (node._icgQueueBusy) return;
      node._icgQueueBusy = true;
      node._icgError = "";
      renderToolbar(node);
      if (!await queueFromHere()) {
        node._icgQueueBusy = false;
        node._icgError = "Could not queue the workflow";
        renderToolbar(node);
      }
    });
    run.disabled = node._icgQueueBusy;
    toolbar.appendChild(status);
    toolbar.appendChild(run);
  }
}

function renderGridState(node) {
  const editable = node._icgState === "paused" && !node._icgBusy;
  for (const cell of node._icgCells.values()) cell.disabled = !editable;
  updateCells(node);
}

function toggleImage(node, index) {
  if (node._icgState !== "paused" || node._icgBusy) return;
  const next = new Set(node._icgSelected);
  if (next.has(index)) next.delete(index);
  else next.add(index);
  select(node, next);
}

function renderGrid(node) {
  const { grid } = node._icg;
  grid.innerHTML = "";
  node._icgCells = new Map();

  for (const item of node._icgItems) {
    const cell = document.createElement("button");
    cell.type = "button";
    cell.className = "icg-cell";
    cell.title = `Image ${item.index + 1} — ${item.width}×${item.height}`;
    cell.setAttribute("aria-label", `Select image ${item.index + 1}`);
    cell.onclick = () => toggleImage(node, item.index);

    const image = document.createElement("img");
    image.className = "icg-thumb";
    image.alt = `Batch image ${item.index + 1}`;
    image.draggable = false;
    image.src = `data:image/jpeg;base64,${item.image}`;

    const badge = document.createElement("span");
    badge.className = "icg-badge";
    badge.textContent = String(item.index + 1);

    const dimensions = document.createElement("span");
    dimensions.className = "icg-dimensions";
    dimensions.textContent = `${item.width}×${item.height}`;

    const check = document.createElement("span");
    check.className = "icg-check";

    cell.appendChild(image);
    cell.appendChild(badge);
    cell.appendChild(dimensions);
    cell.appendChild(check);
    grid.appendChild(cell);
    node._icgCells.set(item.index, cell);
  }
  renderGridState(node);
}

function showBatch(node, waitId, token, items) {
  node._icgRevision += 1;
  node._icgWaitId = waitId;
  node._icgToken = token;
  node._icgItems = Array.isArray(items) ? items : [];
  node._icgSelected = new Set();
  node._icgBusy = false;
  node._icgError = "";
  node._icgStopped = false;
  node._icgResolvedCount = 0;
  node._icgQueueBusy = false;
  node._icgState = "paused";
  renderGrid(node);
  renderToolbar(node);
  resizeChooser(node);
}

function injectStyles() {
  if (document.getElementById("icg-styles")) return;
  const css = `
  .icg-wrap { display:flex; flex-direction:column; gap:7px; box-sizing:border-box;
              height:100%; min-height:0; }
  .icg-grid { flex:1 1 auto; min-height:0; overflow:auto; display:grid;
              grid-template-columns:repeat(auto-fill,minmax(108px,1fr));
              grid-auto-rows:max-content; gap:7px; align-content:start; padding:2px; }
  .icg-cell { position:relative; aspect-ratio:1; min-width:0; overflow:hidden; padding:0;
              border:2px solid transparent; border-radius:5px; background:#1d1d1d;
              cursor:pointer; transition:border-color .1s, box-shadow .1s; }
  .icg-cell:hover:not(:disabled) { border-color:#777; }
  .icg-cell:disabled { cursor:default; opacity:1; }
  .icg-cell.icg-selected { border-color:#55d78a; box-shadow:0 0 0 1px #55d78a; }
  .icg-thumb { position:absolute; inset:0; width:100%; height:100%; object-fit:contain;
               background:#181818; pointer-events:none; }
  .icg-badge, .icg-dimensions, .icg-check { position:absolute; color:#fff;
               background:rgba(0,0,0,.72); border-radius:3px; pointer-events:none; }
  .icg-badge { top:4px; left:4px; min-width:18px; padding:1px 4px; font-size:11px;
               text-align:center; }
  .icg-dimensions { left:4px; bottom:4px; padding:1px 4px; font-size:9px; opacity:.78; }
  .icg-check { top:4px; right:4px; width:20px; height:20px; line-height:20px;
               text-align:center; font-size:14px; font-weight:bold; }
  .icg-cell.icg-selected .icg-check { background:rgba(35,145,80,.92); }
  .icg-toolbar { flex:0 0 auto; display:flex; flex-wrap:wrap; gap:6px; align-items:center; }
  .icg-toolbar button { font-size:12px; padding:4px 10px; cursor:pointer; border-radius:3px;
                        border:1px solid #555; color:#fff; }
  .icg-toolbar button:disabled { opacity:.42; cursor:default; }
  .icg-secondary { background:rgba(45,45,45,.95); }
  .icg-pass { background:rgba(40,130,70,.95); }
  .icg-pass:hover:not(:disabled) { background:rgba(55,160,90,.98); }
  .icg-run { background:rgba(40,90,140,.95); }
  .icg-stop { background:rgba(160,40,40,.9); margin-left:auto; }
  .icg-status { font-size:11px; opacity:.72; padding:0 3px; }
  .icg-error { color:#ff8a8a; opacity:1; }
  .icg-resolved { color:#83e5aa; opacity:.95; }
  `;
  const style = document.createElement("style");
  style.id = "icg-styles";
  style.textContent = css;
  document.head.appendChild(style);
}

function setupChooserNode(node) {
  injectStyles();
  chooserNodes.add(node);

  // This node owns its preview UI; suppress ComfyUI's output-image preview.
  try {
    Object.defineProperty(node, "imgs", {
      configurable: true,
      get() { return undefined; },
      set() { /* suppress */ },
    });
  } catch (e) { /* best effort */ }

  const wrap = document.createElement("div");
  wrap.className = "icg-wrap";
  const grid = document.createElement("div");
  grid.className = "icg-grid";
  const toolbar = document.createElement("div");
  toolbar.className = "icg-toolbar";
  wrap.appendChild(grid);
  wrap.appendChild(toolbar);

  node._icg = { wrap, grid, toolbar };
  node._icgItems = [];
  node._icgCells = new Map();
  node._icgSelected = new Set();
  node._icgState = "idle";
  node._icgBusy = false;
  node._icgRevision = 0;
  node._icgWaitId = null;
  node._icgToken = null;
  node._icgQueueBusy = false;

  node._icgWidget = node.addDOMWidget("image_chooser", "div", wrap, {
    serialize: false,
    getMinHeight: () => widgetFloor(node),
  });

  const onResize = node.onResize;
  node.onResize = function () {
    const result = onResize?.apply(this, arguments);
    syncWidgetWidth(node);
    return result;
  };

  // Removing an actively waiting chooser must not strand the executor.
  const onRemoved = node.onRemoved;
  node.onRemoved = function () {
    if (node._icgState === "paused" && node._icgToken) {
      postChooser(node, { action: "cancel" }).catch(() => {});
    }
    chooserNodes.delete(node);
    return onRemoved?.apply(this, arguments);
  };

  node.setSize([Math.max(node.size?.[0] || 0, MIN_W), node.computeSize()[1]]);
  resizeChooser(node);
}

app.registerExtension({
  name: "datasete.gates.imagechooser",

  setup() {
    api.addEventListener("datasete-image-chooser-show", (event) => {
      const data = event.detail || {};
      const displayId = data.display_id ?? data.id;
      const node = chooserByLocator(displayId);
      if (!node || node.type !== NODE || !node._icg) return;
      showBatch(node, data.id, data.token, data.images);
    });

    // Collapse the previous review as a new execution starts; a fresh chooser
    // event will reopen it if this run reaches the node. An interrupt leaves an
    // honest stopped state.
    api.addEventListener("execution_start", () => {
      for (const node of chooserNodes) {
        if (node.type !== NODE || node._icgState === "idle") continue;
        node._icgRevision += 1;
        node._icgState = "idle";
        node._icgBusy = false;
        node._icgQueueBusy = false;
        node._icgWaitId = null;
        node._icgToken = null;
        resizeChooser(node);
      }
    });
    api.addEventListener("execution_interrupted", () => {
      for (const node of chooserNodes) {
        if (node.type !== NODE
            || (node._icgState !== "paused" && !node._icgQueueBusy)) continue;
        node._icgRevision += 1;
        node._icgState = "resolved";
        node._icgBusy = false;
        node._icgQueueBusy = false;
        node._icgStopped = true;
        node._icgWaitId = null;
        node._icgToken = null;
        renderGridState(node);
        renderToolbar(node);
      }
    });
    api.addEventListener("execution_error", () => {
      for (const node of chooserNodes) {
        if (node.type !== NODE
            || (node._icgState !== "paused" && !node._icgQueueBusy)) continue;
        node._icgRevision += 1;
        node._icgState = "resolved";
        node._icgBusy = false;
        node._icgQueueBusy = false;
        node._icgStopped = true;
        node._icgError = "Execution stopped with an error";
        node._icgWaitId = null;
        node._icgToken = null;
        renderGridState(node);
        renderToolbar(node);
      }
    });
  },

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

    const onNodeCreated = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
      const result = onNodeCreated?.apply(this, arguments);
      setupChooserNode(this);
      return result;
    };

    const onConfigure = nodeType.prototype.onConfigure;
    nodeType.prototype.onConfigure = function () {
      const result = onConfigure?.apply(this, arguments);
      if (this._icg) syncWidgetWidth(this);
      return result;
    };
  },
});