import { app } from "../../../scripts/app.js";
import { allow_debug } from "./js_shared.js";
import { BaseSmartWidget, BaseSmartWidgetManager } from "./makadi/BaseSmartWidget.js";
import { SmartButton } from "./makadi/SmartButton.js";
import { Shapes } from "./utils.js";

app.registerExtension({
  name: "iTools.previewNode",
  async beforeRegisterNodeDef(nodeType, nodeData, app) {
    if (nodeData.name === "iToolsPreviewImage") {
    }
  },
  async nodeCreated(node) {
    if (node.comfyClass !== "iToolsPreviewImage") {
      return;
    }

    // init size
    node.size = [350, 350];

    if (allow_debug) console.log("node", node);

    // vars
    let a = null;
    let b = null;
    let c = null;
    let noteButton = null;
    let drawNoteButton = () => {};
    let compare = false;
    let imgsBeforeCompare = null;
    let imagesTracked = [];
    const MAX_IMAGES = 8;

    let mouse = {
      mouseInNode: false,
      x: 0,
      y: 0,
    };

    // other vars
    let toastShownCountH = 0;
    let toastShownCountPC = 0;
    let toastShownCountI = 0;
    const MAX_TOAST_SHOWS = 2;

    function pushToImgs(newImage) {
      // Check if image is undefined or null
      if (!newImage || !newImage.naturalWidth) {
        if (allow_debug) console.log("Undefined or null image, skipping");
        return imagesTracked;
      }

      const imageExists = imagesTracked.some((img) => {
        // Extract filename from URLs by removing the random parameter
        const getFilename = (url) => {
          if (!url) return "";
          const match = url.match(/filename=([^&]+)/);
          return match ? match[1] : "";
        };

        const newImageFilename = getFilename(newImage.src);
        const existingImageFilename = getFilename(img.src);

        return newImageFilename === existingImageFilename;
      });

      if (imageExists) {
        if (allow_debug) console.log("Image already exists, skipping");
        return imagesTracked;
      }

      imagesTracked.push(newImage);
      // If array exceeds MAX_IMAGES, remove oldest image(s)
      if (imagesTracked.length > MAX_IMAGES) {
        imagesTracked.shift(); // removes first (oldest) element
      }
      if (allow_debug) console.log("imagesTracked list updated, length:", imagesTracked.length);
      return imagesTracked;
    }

    function cycleImgs() {
      if (imagesTracked.length === 1) {
        app.extensionManager.toast.add({
          severity: "info",
          summary: "iTools!",
          detail: "Only this image exist in this node history",
          life: 2000,
        });
        return;
      }

      // Switch to previous imagesTracked if available
      if (node.imgs && imagesTracked.length > 1) {
        node.imgs = imagesTracked; // Show previous image
      }
    }

    function togglingLastTwoImages() {
      if (allow_debug) console.log("imagesTracked", imagesTracked);
      // Only cycle between last two imagesTracked
      if (imagesTracked.length > 1) {
        node.imageIndex = 0;

        // Get current image
        const currentImg = node.imgs[0];

        // Get last two imagesTracked from history
        const lastTwo = imagesTracked.slice(-2);

        // If current is last image, show second to last
        // If current is second to last, show last
        const nextImg = currentImg === lastTwo[1] ? lastTwo[0] : lastTwo[1];

        // Update display
        node.imgs = [nextImg];

        // Update button text with underline using Unicode
        const isShowingCurrent = nextImg === lastTwo[1];
        b.text = isShowingCurrent ? "[Current] | Previous" : "Current | [Previous]";

        if (allow_debug) console.log("Toggling between last two imagesTracked");
      } else {
        if (toastShownCountPC < MAX_TOAST_SHOWS) {
          app.extensionManager.toast.add({
            severity: "info",
            summary: "iTools!",
            detail: "You must execute this node at least twice",
            life: 2000,
          });
          toastShownCountPC++;
        }
      }
    }

    function toggleButtonActivation() {
      c.isActive = compare;
      if (!compare) {
        c.color = c.originalColor;
        c.textColor = c.originalTextColor;
      } else {
        c.color = "#80a1c0";
        c.textColor = "black";
      }
    }

    function showButtons() {
      a.isVisible = true;
      b.isVisible = true;
      // c.isVisible = true;
    }

    function createButtons(startVisible = true) {
      a = new SmartButton(75, 8, 55, 20, node, "History");
      a.allowVisualHover = true;
      a.textYoffset = -0;
      a.isVisible = startVisible;
      a.shape = Shapes.ROUND_L;
      a.roundRadius = 5;
      a.outlineWidth = 1;
      a.outlineColor = "#656565";
      a.color = "#222222";
      a.font = "12px Arial";
      a.onClick = () => {
        if (compare) {
          // cancel compare
          compare = !compare;
          applyCompareState();
          toggleButtonActivation(c, compare);
        }
        if (!node.imgs) {
          if (toastShownCountH < MAX_TOAST_SHOWS) {
            app.extensionManager.toast.add({
              severity: "info",
              summary: "iTools!",
              detail: "No images in this node history",
              life: 2000,
            });
            toastShownCountH++;
          }
          return;
        }
        if (node.imgs.length > 1) {
          togglingLastTwoImages();
        } else {
          cycleImgs();
        }
      };

      b = new SmartButton(75 + 55, 8, 120, 20, node, "[Current] | Previous");
      b.allowVisualHover = true;
      b.textYoffset = -0;
      b.isVisible = startVisible;
      b.shape = Shapes.ROUND_R;
      b.roundRadius = 5;
      b.outlineWidth = 1;
      b.outlineColor = "#656565";
      b.color = "#222222";
      b.font = "12px Arial";
      b.onClick = () => {
        if (compare) {
          // cancel compare
          compare = !compare;
          applyCompareState();
          toggleButtonActivation(c, compare);
        }
        togglingLastTwoImages();
      };

      c = new SmartButton(75 + 55 + 125, 8 + 1, 18, 20, node, "|");
      c.allowVisualHover = true;
      c.textYoffset = -0.05;
      c.isVisible = startVisible;
      c.shape = Shapes.CIRCLE;
      //c.roundRadius = 5;
      c.outlineWidth = 1;
      c.outlineColor = "#656565";
      c.color = "#222222";
      c.activeColor = c.font = "12px Arial";
      c.onClick = () => {
        // reset togglingLastTwoImages
        if (b.text !== "[Current] | Previous") togglingLastTwoImages();

        // start compare
        if (imagesTracked.length <= 1) {
          if (toastShownCountI < MAX_TOAST_SHOWS) {
            app.extensionManager.toast.add({
              severity: "info",
              summary: "iTools!",
              detail: "You must execute this node at least twice",
              life: 2000,
            });
            toastShownCountI++;
          }
          return;
        }
        compare = !compare;
        applyCompareState();
        toggleButtonActivation(c, compare);
      };

      // Sits on the top left of the preview area
      noteButton = new SmartButton(8, 40, 54, 17, node, "Add Note");
      noteButton.allowVisualHover = true;
      noteButton.textYoffset = -0;
      noteButton.isVisible = false;
      noteButton.shape = Shapes.ROUND;
      noteButton.roundRadius = 5;
      noteButton.outlineWidth = 1;
      noteButton.outlineColor = "#656565";
      noteButton.color = "#222222AA";
      noteButton.font = "11px Arial";
      noteButton.onClick = () => {
        const img = noteTargetImage();
        if (!img) return;
        app.extensionManager.dialog
          .prompt({
            title: "iTools Preview Image",
            message: "Note for this image (leave empty to remove)",
            defaultValue: getNote(img), // key name differs between frontend versions
            default: getNote(img),
          })
          .then((value) => {
            if (value === null || value === undefined) return; // cancelled
            setNote(img, value.trim());
            node.setDirtyCanvas(true, true);
          });
      };

      // The frontend paints the image from a deferred microtask, which lands on
      // top of anything the widget pass drew, so this button is painted from that
      // same deferred pass instead (see drawNoteLayer). Clicks still use isVisible.
      const origButtonDraw = noteButton.draw;
      noteButton.draw = () => {};
      drawNoteButton = (ctx) => {
        if (noteButton.isVisible) origButtonDraw.call(noteButton, ctx);
      };
    }

    // Notes are stored per image filename in node.properties, so they are saved
    // with the workflow and follow the image through history / compare instead of
    // being baked into the image itself
    function getNotes() {
      if (!node.properties.iToolsImageNotes) node.properties.iToolsImageNotes = {};
      return node.properties.iToolsImageNotes;
    }

    function getNote(img) {
      const key = getImageKey(img);
      return (key && getNotes()[key]) || "";
    }

    function setNote(img, text) {
      const key = getImageKey(img);
      if (!key) return;
      if (text) getNotes()[key] = text;
      else delete getNotes()[key];
    }

    // The image a new note applies to: the one currently on screen, or the
    // current (right hand) image while comparing
    function noteTargetImage() {
      if (!node.imgs?.length) return null;
      if (compare && node.imgs.length > 1) return node.imgs.at(-1);
      // In the grid state overIndex is the image the mouse is over
      if (node.imageIndex == null && node.imgs.length > 1) return node.imgs[node.overIndex ?? 0];
      return node.imgs[node.imageIndex ?? 0] || node.imgs[0];
    }

    // Mirrors the frontend renderPreview geometry so the note sits exactly on the
    // image in both preview states: single image and the multi image grid
    function collectNoteTargets(widget_width, shiftY, computedHeight) {
      if (!node.imgs?.length) return [];

      // Grid state: the frontend fills node.imageRects with the cell of each image
      if (node.imageIndex == null && node.imgs.length > 1) {
        const cells = node.imageRects;
        if (!cells?.length) return [];
        return node.imgs
          .map((img, i) => {
            const cell = cells[i];
            if (!cell || !img?.width) return null;
            const [cellX, cellY, cellW, cellH] = cell;
            const ratio = Math.min(cellW / img.width, cellH / img.height);
            const w = img.width * ratio;
            const h = img.height * ratio;
            return { img, rect: { x: cellX + (cellW - w) / 2, y: cellY + (cellH - h) / 2, w, h } };
          })
          .filter(Boolean);
      }

      // Single image state, scale is capped at 1 (no upscaling) like the frontend
      const img = node.imgs[node.imageIndex ?? 0];
      if (!img?.naturalWidth) return [];
      const sizeTextHeight = app.extensionManager.setting.get("Comfy.Node.AllowImageSizeDraw") ? 15 : 0;
      const dw = widget_width;
      const dh = (computedHeight ?? node.size[1] - shiftY) - sizeTextHeight;
      const scale = Math.min(dw / img.naturalWidth, dh / img.naturalHeight, 1);
      const w = img.naturalWidth * scale;
      const h = img.naturalHeight * scale;
      return [{ img, rect: { x: (dw - w) / 2, y: (dh - h) / 2 + shiftY, w, h } }];
    }

    // The frontend defers its drawImage calls to a microtask, so queue the notes
    // and the button right after to land on top of the images instead of under them
    function drawNoteLayer(ctx, widget_width, shiftY, computedHeight) {
      const targets = collectNoteTargets(widget_width, shiftY || 0, computedHeight).filter((t) =>
        getNote(t.img)
      );

      const transform = ctx.getTransform();
      queueMicrotask(() => {
        ctx.save();
        ctx.setTransform(transform);
        for (const target of targets) drawNoteBar(ctx, target.rect, getNote(target.img));
        drawNoteButton(ctx);
        ctx.restore();
      });
    }

    // iToolsPreviewImage gets one image per run, while iToolsCompareImage gets two.
    // Feeding node.imgs the last two tracked images puts this node in the same
    // two-image preview state compare_node draws from.
    function applyCompareState() {
      if (compare) {
        const lastTwo = imagesTracked.slice(-2);
        imgsBeforeCompare = node.imgs;
        node.imgs = [lastTwo[0], lastTwo[1]]; // A = previous, B = current
        node.imageIndex = null; // multi image view, same as compare node
      } else {
        node.imgs = imgsBeforeCompare || [imagesTracked.at(-1)];
        node.imageIndex = 0;
        imgsBeforeCompare = null;
      }
      patchPreviewWidget(); // widget may already be swapped by the assignment above
      node.setDirtyCanvas(true, true);
    }

    // Patches the image preview widget draw. Re-applied on demand because
    // assigning node.imgs makes the frontend rebuild the widget, dropping the patch.
    function patchPreviewWidget() {
      // ImagePreviewWidget is added by the frontend under this name
      const previewWidget =
        node.widgets?.find((widget) => widget.name === "$$canvas-image-preview") ||
        node.widgets?.find((widget) => !(widget instanceof BaseSmartWidget) && widget.drawWidget);

      // Pinned to the node, top left of the preview area
      if (noteButton) {
        const targetImg = noteTargetImage();
        // Hidden in the history grid, there is no single image to attach a note to
        const inHistoryGrid = !compare && node.imageIndex == null && node.imgs?.length > 1;
        noteButton.isVisible = !!previewWidget && !!node.imgs?.length && !inHistoryGrid;
        if (typeof previewWidget?.y === "number") noteButton.myY = previewWidget.y + 6;
        noteButton.text = getNote(targetImg) ? "Edit Note" : "Add Note";
      }

      if (!previewWidget) return;

      const comparing = () => compare && node.imgs?.length > 1;

      // The canvas uses widget.draw when it exists, otherwise widget.drawWidget,
      // so both entry points need the override
      // Only wrap draw if the widget already has one, adding it would change dispatch
      if (typeof previewWidget.draw === "function" && previewWidget.draw !== previewWidget._itoolsDrawWrapper) {
        const originalDraw = previewWidget.draw;
        const drawWrapper = function (ctx, node, widget_width, y, widget_height, lowQuality) {
          if (comparing()) {
            drawImgOverlay(mouse, node, widget_width, y, ctx, compare, getNote);
            drawNoteButton(ctx); // compare draws its images synchronously
          } else {
            originalDraw.call(this, ctx, node, widget_width, y, widget_height, lowQuality);
            drawNoteLayer(ctx, widget_width, y, this.computedHeight ?? widget_height);
          }
        };
        previewWidget._itoolsDrawWrapper = drawWrapper;
        previewWidget.draw = drawWrapper;
      }

      if (previewWidget.drawWidget !== previewWidget._itoolsDrawWidgetWrapper) {
        const originalDrawWidget = previewWidget.drawWidget;
        const drawWidgetWrapper = function (ctx, options) {
          const width = options?.width ?? node.size[0];
          if (comparing()) {
            drawImgOverlay(mouse, node, width, this.y, ctx, compare, getNote);
            drawNoteButton(ctx); // compare draws its images synchronously
          } else {
            originalDrawWidget?.call(this, ctx, options);
            drawNoteLayer(ctx, width, this.y, this.computedHeight);
          }
        };
        previewWidget._itoolsDrawWidgetWrapper = drawWidgetWrapper;
        previewWidget.drawWidget = drawWidgetWrapper;
      }
    }

    // Canvas level hook, same approach as BaseSmartWidget: runs every frame so the
    // patch is re-applied as soon as the frontend rebuilds the preview widget
    const origCanvasDrawForeground = app.canvas.onDrawForeground;
    app.canvas.onDrawForeground = (ctx) => {
      if (origCanvasDrawForeground) origCanvasDrawForeground.call(app.canvas, ctx);
      if (node.graph) patchPreviewWidget();
    };

    createButtons();

    node.onExecuted = async function (message) {
      // Wait for image to be loaded
      for (let i = 0; i < 20 && !node.imgs; i++) {
        if (allow_debug) console.log("wait..", i);
        await new Promise((resolve) => setTimeout(resolve, 100));
      }
      if (!node.imgs) return;

      showButtons();

      // Reset togglingLastTwoImages
      if (b.text !== "[Current] | Previous") togglingLastTwoImages();
      node.setDirtyCanvas(true, false);

      setTimeout(() => {
        // push last image
        const lastImage = node.imgs?.at(-1);
        if (!imagesTracked.some((img) => img === lastImage)) {
          pushToImgs(lastImage);
        }

        // Keep comparing against the freshly generated image
        if (compare) applyCompareState();

        // Override draw function in ImagePreviewWidget
        patchPreviewWidget();
      }, 300);
    };

    node.onResize = function (newSize) {
      // limit width size while resizing
      node.size[0] = Math.max(285, newSize[0]);
      // if (allow_debug) console.log("Node", node);
    };

    const m = new BaseSmartWidgetManager(node, "iToolsPreviewImage");
    const origOnRemoved = node.onRemoved;
    node.onRemoved = function () {
      origOnRemoved?.apply(this, arguments);
      m.destroy();
    };

    node.onMouseEnter = (e) => {
      mouse.mouseInNode = true;
    };

    node.onMouseLeave = (e) => {
      mouse.mouseInNode = false;
    };

    node.onMouseMove = (e, pos) => {
      if (mouse.mouseInNode) {
        const graphMouse = app.canvas.graph_mouse;
        mouse.x = graphMouse[0] - node.pos[0];
        mouse.y = graphMouse[1] - node.pos[1];
      }
    };
  },
  
});

const compareWay = app.extensionManager.setting.get("iTools.Nodes.Compare Mode", "makadi");

// Notes key off the output filename so the same image keeps its note across
// runs, history cycling and compare
function getImageKey(img) {
  if (!img?.src) return "";
  const match = img.src.match(/filename=([^&]+)/);
  return match ? decodeURIComponent(match[1]) : img.src;
}

// Note bar sitting on the bottom of the image, clipped to it
function drawNoteBar(ctx, rect, text) {
  if (!text || !rect) return;
  const fontSize = Math.max(8, Math.min(12, Math.round(rect.w / 16)));
  const padding = Math.round(fontSize / 2);
  const barHeight = fontSize + padding * 2;
  const barY = rect.y + rect.h - barHeight;

  ctx.save();
  ctx.beginPath();
  ctx.rect(rect.x, barY, rect.w, barHeight);
  ctx.clip();
  ctx.fillStyle = "#000000AA";
  ctx.fillRect(rect.x, barY, rect.w, barHeight);
  ctx.fillStyle = "white";
  ctx.font = `${fontSize}px monospace`;
  ctx.textAlign = "left";
  ctx.textBaseline = "middle";
  ctx.fillText(text, rect.x + padding, barY + barHeight / 2);
  ctx.restore();
}

function drawImgOverlay(mouse, node, widget_width, y, ctx, compareMode = false, getNote = null) {
  if (!compareMode || !node.imgs || node.imgs.length < 2) return;
  y = y ? y : 0; // Ensure y is defined

  const img1 = node.imgs[0]; // previous
  const img2 = node.imgs[1]; // current
  if (!img1 || !img2) {
    if (allow_debug) console.log("No previous image to compare with");
    return;
  }

  const dw = widget_width;
  const dh = node.size[1] - y;

  // Force both images to the same height (dh), scaling down if too wide
  const getParams = (img) => {
    const scale = dh / img.naturalHeight;
    const w = img.naturalWidth * scale;
    const finalScale = w > dw ? dw / w : 1;

    const finalW = w * finalScale;
    const finalH = dh * finalScale;

    return {
      x: (dw - finalW) / 2,
      y: y + (dh - finalH) / 2,
      w: finalW,
      h: finalH,
    };
  };

  const p1 = getParams(img1);
  const p2 = getParams(img2);

  // Shared interaction bounds (the container area)
  const viewW = Math.max(p1.w, p2.w);
  const imgX = (dw - viewW) / 2;

  let mouseX;
  if (compareWay === "makadi") {
    const graphMouse = app.canvas.graph_mouse;
    mouseX = graphMouse[0] - node.pos[0];
  } else {
    mouseX = mouse.mouseInNode ? mouse.x : dw / 2;
  }

  const splitX = Math.max(imgX, Math.min(mouseX, imgX + viewW));

  const left = compareWay === "makadi" ? { img: img1, p: p1 } : { img: img2, p: p2 };
  const right = compareWay === "makadi" ? { img: img2, p: p2 } : { img: img1, p: p1 };

  // Draw Left Side
  ctx.save();
  ctx.beginPath();
  ctx.rect(0, y, splitX, dh);
  ctx.clip();
  ctx.drawImage(left.img, left.p.x, left.p.y, left.p.w, left.p.h);
  drawNoteBar(ctx, left.p, getNote?.(left.img)); // clipped to this half
  ctx.restore();

  // Draw Right Side
  ctx.save();
  ctx.beginPath();
  ctx.rect(splitX, y, dw - splitX, dh);
  ctx.clip();
  ctx.drawImage(right.img, right.p.x, right.p.y, right.p.w, right.p.h);
  drawNoteBar(ctx, right.p, getNote?.(right.img));
  ctx.restore();
}