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

function chainCallback(object, property, callback) {
  const original = object[property];
  if (original) {
    object[property] = function (...args) {
      const result = original.apply(this, args);
      callback.apply(this, args);
      return result;
    };
  } else {
    object[property] = callback;
  }
}

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

function clamp(value, min, max) {
  return Math.max(min, Math.min(max, value));
}

function parseRegions(value) {
  try {
    const parsed = JSON.parse(value || "[]");
    if (Array.isArray(parsed)) return parsed;
    if (parsed && Array.isArray(parsed.regions)) return parsed.regions;
    if (parsed && typeof parsed === "object") return [parsed];
  } catch (error) {
    console.warn("UR Region Editor: failed to parse regions_json", error);
  }
  return [];
}

function normalizeRect(region) {
  const x = Number(region.x ?? region.x1 ?? 0);
  const y = Number(region.y ?? region.y1 ?? 0);
  if (region.w != null || region.h != null) {
    return {
      x,
      y,
      w: Number(region.w ?? 0),
      h: Number(region.h ?? 0),
    };
  }
  const x2 = Number(region.x2 ?? region.right ?? x);
  const y2 = Number(region.y2 ?? region.bottom ?? y);
  return {
    x: Math.min(x, x2),
    y: Math.min(y, y2),
    w: Math.abs(x2 - x),
    h: Math.abs(y2 - y),
  };
}

function createButton(label, title) {
  const button = document.createElement("button");
  button.textContent = label;
  button.title = title || label;
  button.style.cssText = "height:24px;padding:0 8px;border:1px solid #555;background:#2b2b2b;color:#ddd;border-radius:4px;cursor:pointer;font-size:12px;";
  button.addEventListener("mouseenter", () => {
    button.style.background = "#3a3a3a";
  });
  button.addEventListener("mouseleave", () => {
    button.style.background = button.dataset.active === "true" ? "#4b5f7a" : "#2b2b2b";
  });
  return button;
}

function createNumberInput(value, min, title) {
  const input = document.createElement("input");
  input.type = "number";
  input.value = String(value);
  input.min = String(min ?? 0);
  input.title = title || "";
  input.style.cssText = "height:22px;width:76px;background:#111;color:#ddd;border:1px solid #555;border-radius:4px;padding:0 4px;font-size:12px;";
  return input;
}

class RegionEditorWidget {
  constructor(node) {
    this.node = node;
    this.regionsWidget = findWidget(node, "regions_json");
    this.widthWidget = findWidget(node, "canvas_width");
    this.heightWidget = findWidget(node, "canvas_height");
    this.defaultFrameWidget = findWidget(node, "default_frame_index");

    this.regions = parseRegions(this.regionsWidget?.value);
    this.mode = "rect";
    this.activeFrame = Number(this.defaultFrameWidget?.value ?? 0);
    this.selectedIndex = -1;
    this.dragStart = null;
    this.currentDraft = null;
    this.polygonDraft = [];
    this.background = null;
    this.video = null;
    this.videoFps = 24;

    this.root = document.createElement("div");
    this.root.style.cssText = "display:flex;flex-direction:column;gap:6px;width:100%;height:100%;box-sizing:border-box;background:#121212;border:1px solid #303030;border-radius:6px;padding:6px;overflow:hidden;";

    this.toolbar = document.createElement("div");
    this.toolbar.style.cssText = "display:flex;align-items:center;gap:6px;flex-wrap:wrap;";

    this.rectBtn = createButton("Rect", "Draw rectangle by dragging");
    this.circleBtn = createButton("Circle", "Draw circle by dragging");
    this.polyBtn = createButton("Poly", "Click points, then Finish Poly");
    this.finishPolyBtn = createButton("Finish Poly", "Finish current polygon");
    this.undoPointBtn = createButton("Undo Point", "Remove last polygon point");
    this.deleteBtn = createButton("Delete", "Delete selected region");
    this.clearBtn = createButton("Clear", "Clear all regions");
    this.loadBgBtn = createButton("Load BG", "Load local image/video as editor background only");

    this.frameInput = createNumberInput(this.activeFrame, 0, "Current keyframe index");
    this.frameLabel = document.createElement("span");
    this.frameLabel.textContent = "Frame";
    this.frameLabel.style.cssText = "color:#aaa;font-size:12px;";

    this.fileInput = document.createElement("input");
    this.fileInput.type = "file";
    this.fileInput.accept = "image/*,video/*";
    this.fileInput.style.display = "none";

    this.videoControls = document.createElement("div");
    this.videoControls.style.cssText = "display:none;align-items:center;gap:6px;width:100%;";
    this.playBtn = createButton("Play", "Play/pause background video");
    this.videoSlider = document.createElement("input");
    this.videoSlider.type = "range";
    this.videoSlider.min = "0";
    this.videoSlider.max = "0";
    this.videoSlider.step = "0.001";
    this.videoSlider.value = "0";
    this.videoSlider.style.cssText = "flex:1;";
    this.fpsInput = createNumberInput(this.videoFps, 1, "FPS used to convert video time to frame index");
    this.useVideoFrameBtn = createButton("Use Time", "Set keyframe from current video time");

    this.canvas = document.createElement("canvas");
    this.canvas.width = Number(this.widthWidget?.value ?? 768);
    this.canvas.height = Number(this.heightWidget?.value ?? 512);
    this.canvas.style.cssText = "width:100%;height:auto;max-height:520px;object-fit:contain;background:#181818;border:1px solid #333;border-radius:4px;cursor:crosshair;";
    this.ctx = this.canvas.getContext("2d");

    this.status = document.createElement("div");
    this.status.style.cssText = "color:#aaa;font-size:12px;white-space:pre-wrap;line-height:1.3;";

    this.toolbar.append(
      this.rectBtn,
      this.circleBtn,
      this.polyBtn,
      this.finishPolyBtn,
      this.undoPointBtn,
      this.deleteBtn,
      this.clearBtn,
      this.loadBgBtn,
      this.frameLabel,
      this.frameInput,
      this.fileInput,
    );
    this.videoControls.append(this.playBtn, this.videoSlider, document.createTextNode("FPS"), this.fpsInput, this.useVideoFrameBtn);
    this.root.append(this.toolbar, this.videoControls, this.canvas, this.status);

    this.bindEvents();
    this.setMode("rect");
    this.syncWidgets(false);
    this.render();
  }

  bindEvents() {
    this.rectBtn.onclick = () => this.setMode("rect");
    this.circleBtn.onclick = () => this.setMode("circle");
    this.polyBtn.onclick = () => this.setMode("polygon");
    this.finishPolyBtn.onclick = () => this.finishPolygon();
    this.undoPointBtn.onclick = () => {
      this.polygonDraft.pop();
      this.render();
    };
    this.deleteBtn.onclick = () => this.deleteSelected();
    this.clearBtn.onclick = () => {
      this.regions = [];
      this.selectedIndex = -1;
      this.polygonDraft = [];
      this.syncWidgets(true);
      this.render();
    };
    this.loadBgBtn.onclick = () => this.fileInput.click();

    this.frameInput.addEventListener("change", () => {
      this.activeFrame = Math.max(0, Number(this.frameInput.value || 0));
      if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame;
      this.render();
    });
    this.fpsInput.addEventListener("change", () => {
      this.videoFps = Math.max(1, Number(this.fpsInput.value || 24));
    });

    this.fileInput.addEventListener("change", () => this.loadBackgroundFile());
    this.playBtn.onclick = () => this.toggleVideo();
    this.videoSlider.addEventListener("input", () => {
      if (!this.video) return;
      this.video.currentTime = Number(this.videoSlider.value || 0);
      this.drawVideoFrame();
    });
    this.useVideoFrameBtn.onclick = () => {
      if (!this.video) return;
      this.activeFrame = Math.max(0, Math.round(this.video.currentTime * this.videoFps));
      this.frameInput.value = String(this.activeFrame);
      if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame;
      this.render();
    };

    this.canvas.addEventListener("mousedown", (event) => this.onMouseDown(event));
    this.canvas.addEventListener("mousemove", (event) => this.onMouseMove(event));
    this.canvas.addEventListener("mouseup", (event) => this.onMouseUp(event));
    this.canvas.addEventListener("dblclick", () => this.finishPolygon());
    this.canvas.addEventListener("contextmenu", (event) => {
      event.preventDefault();
      const point = this.toCanvasPoint(event);
      const index = this.hitTest(point.x, point.y);
      if (index >= 0) {
        this.selectedIndex = index;
        this.deleteSelected();
      }
    });
    document.addEventListener("keydown", (event) => {
      if (!this.root.isConnected) return;
      if (event.key === "Delete" || event.key === "Backspace") this.deleteSelected();
      if (event.key === "Enter" && this.mode === "polygon") this.finishPolygon();
      if (event.key === "Escape") {
        this.currentDraft = null;
        this.polygonDraft = [];
        this.render();
      }
    });

    if (this.regionsWidget) {
      const previous = this.regionsWidget.callback;
      this.regionsWidget.callback = (...args) => {
        if (previous) previous.apply(this.regionsWidget, args);
        this.regions = parseRegions(this.regionsWidget.value);
        this.render();
      };
    }
  }

  setMode(mode) {
    this.mode = mode;
    for (const [button, name] of [[this.rectBtn, "rect"], [this.circleBtn, "circle"], [this.polyBtn, "polygon"]]) {
      button.dataset.active = name === mode ? "true" : "false";
      button.style.background = name === mode ? "#4b5f7a" : "#2b2b2b";
    }
    this.currentDraft = null;
    this.render();
  }

  syncWidgets(markDirty) {
    const text = JSON.stringify(this.regions);
    if (this.regionsWidget) this.regionsWidget.value = text;
    if (this.widthWidget) this.widthWidget.value = this.canvas.width;
    if (this.heightWidget) this.heightWidget.value = this.canvas.height;
    if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame;
    if (markDirty) app.graph?.setDirtyCanvas(true, true);
  }

  loadBackgroundFile() {
    const file = this.fileInput.files?.[0];
    if (!file) return;
    const url = URL.createObjectURL(file);
    this.background = null;
    this.video = null;
    this.videoControls.style.display = "none";
    if (file.type.startsWith("video/")) {
      const video = document.createElement("video");
      video.src = url;
      video.muted = true;
      video.playsInline = true;
      video.addEventListener("loadedmetadata", () => {
        this.video = video;
        this.canvas.width = video.videoWidth || this.canvas.width;
        this.canvas.height = video.videoHeight || this.canvas.height;
        this.videoSlider.max = String(video.duration || 0);
        this.videoControls.style.display = "flex";
        this.drawVideoFrame();
        this.syncWidgets(true);
      });
      video.addEventListener("timeupdate", () => {
        this.videoSlider.value = String(video.currentTime);
        this.drawVideoFrame();
      });
      return;
    }

    const image = new Image();
    image.onload = () => {
      this.background = image;
      this.canvas.width = image.width;
      this.canvas.height = image.height;
      this.syncWidgets(true);
      this.render();
    };
    image.src = url;
  }

  toggleVideo() {
    if (!this.video) return;
    if (this.video.paused) {
      this.video.play();
      this.playBtn.textContent = "Pause";
    } else {
      this.video.pause();
      this.playBtn.textContent = "Play";
    }
  }

  drawVideoFrame() {
    this.render();
  }

  toCanvasPoint(event) {
    const rect = this.canvas.getBoundingClientRect();
    return {
      x: clamp((event.clientX - rect.left) * (this.canvas.width / rect.width), 0, this.canvas.width),
      y: clamp((event.clientY - rect.top) * (this.canvas.height / rect.height), 0, this.canvas.height),
    };
  }

  onMouseDown(event) {
    event.preventDefault();
    const point = this.toCanvasPoint(event);
    const hit = this.hitTest(point.x, point.y);
    if (hit >= 0 && event.shiftKey) {
      this.selectedIndex = hit;
      this.render();
      return;
    }
    if (this.mode === "polygon") {
      this.polygonDraft.push(point);
      this.render();
      return;
    }
    this.dragStart = point;
    this.currentDraft = { type: this.mode, frame_index: this.activeFrame, start: point, end: point };
  }

  onMouseMove(event) {
    if (!this.currentDraft) return;
    this.currentDraft.end = this.toCanvasPoint(event);
    this.render();
  }

  onMouseUp(event) {
    if (!this.currentDraft) return;
    this.currentDraft.end = this.toCanvasPoint(event);
    const region = this.draftToRegion(this.currentDraft);
    this.currentDraft = null;
    if (region) {
      this.regions.push(region);
      this.selectedIndex = this.regions.length - 1;
      this.syncWidgets(true);
    }
    this.render();
  }

  draftToRegion(draft) {
    const x1 = draft.start.x;
    const y1 = draft.start.y;
    const x2 = draft.end.x;
    const y2 = draft.end.y;
    if (draft.type === "rect") {
      const x = Math.min(x1, x2);
      const y = Math.min(y1, y2);
      const w = Math.abs(x2 - x1);
      const h = Math.abs(y2 - y1);
      if (w < 3 || h < 3) return null;
      return { type: "rect", frame_index: this.activeFrame, x: Math.round(x), y: Math.round(y), w: Math.round(w), h: Math.round(h) };
    }
    if (draft.type === "circle") {
      const r = Math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2);
      if (r < 3) return null;
      return { type: "circle", frame_index: this.activeFrame, cx: Math.round(x1), cy: Math.round(y1), r: Math.round(r) };
    }
    return null;
  }

  finishPolygon() {
    if (this.polygonDraft.length < 3) return;
    this.regions.push({
      type: "polygon",
      frame_index: this.activeFrame,
      points: this.polygonDraft.map((point) => [Math.round(point.x), Math.round(point.y)]),
    });
    this.selectedIndex = this.regions.length - 1;
    this.polygonDraft = [];
    this.syncWidgets(true);
    this.render();
  }

  deleteSelected() {
    if (this.selectedIndex < 0 || this.selectedIndex >= this.regions.length) return;
    this.regions.splice(this.selectedIndex, 1);
    this.selectedIndex = -1;
    this.syncWidgets(true);
    this.render();
  }

  hitTest(x, y) {
    for (let i = this.regions.length - 1; i >= 0; i--) {
      const region = this.regions[i];
      if (Number(region.frame_index ?? 0) !== this.activeFrame) continue;
      const type = String(region.type || "rect").toLowerCase();
      if (type === "rect") {
        const rect = normalizeRect(region);
        if (x >= rect.x && x <= rect.x + rect.w && y >= rect.y && y <= rect.y + rect.h) return i;
      } else if (type === "circle") {
        const dx = x - Number(region.cx ?? region.x ?? 0);
        const dy = y - Number(region.cy ?? region.y ?? 0);
        if (dx * dx + dy * dy <= Number(region.r ?? region.radius ?? 0) ** 2) return i;
      }
    }
    return -1;
  }

  render() {
    const ctx = this.ctx;
    ctx.clearRect(0, 0, this.canvas.width, this.canvas.height);
    ctx.fillStyle = "#181818";
    ctx.fillRect(0, 0, this.canvas.width, this.canvas.height);
    if (this.video) ctx.drawImage(this.video, 0, 0, this.canvas.width, this.canvas.height);
    else if (this.background) ctx.drawImage(this.background, 0, 0, this.canvas.width, this.canvas.height);
    else this.drawGrid();

    this.regions.forEach((region, index) => this.drawRegion(region, index === this.selectedIndex));
    if (this.currentDraft) this.drawDraft(this.currentDraft);
    if (this.polygonDraft.length) this.drawPolygonDraft();
    this.status.textContent = `Mode: ${this.mode} | Frame: ${this.activeFrame} | Regions: ${this.regions.length}\nLoad BG is editor-only; regions_json is portable and saved in the workflow.`;
  }

  drawGrid() {
    const ctx = this.ctx;
    ctx.strokeStyle = "#2a2a2a";
    ctx.lineWidth = 1;
    for (let x = 0; x <= this.canvas.width; x += 64) {
      ctx.beginPath();
      ctx.moveTo(x, 0);
      ctx.lineTo(x, this.canvas.height);
      ctx.stroke();
    }
    for (let y = 0; y <= this.canvas.height; y += 64) {
      ctx.beginPath();
      ctx.moveTo(0, y);
      ctx.lineTo(this.canvas.width, y);
      ctx.stroke();
    }
    ctx.fillStyle = "#777";
    ctx.font = "16px sans-serif";
    ctx.textAlign = "center";
    ctx.fillText("Load image/video background, then draw regions", this.canvas.width / 2, this.canvas.height / 2);
  }

  drawRegion(region, selected) {
    const ctx = this.ctx;
    const type = String(region.type || "rect").toLowerCase();
    ctx.save();
    ctx.strokeStyle = selected ? "#ffcc00" : "#35c759";
    ctx.fillStyle = selected ? "rgba(255,204,0,0.20)" : "rgba(53,199,89,0.20)";
    ctx.lineWidth = selected ? 3 : 2;
    if (type === "rect") {
      const rect = normalizeRect(region);
      ctx.fillRect(rect.x, rect.y, rect.w, rect.h);
      ctx.strokeRect(rect.x, rect.y, rect.w, rect.h);
    } else if (type === "circle") {
      ctx.beginPath();
      ctx.arc(Number(region.cx ?? region.x ?? 0), Number(region.cy ?? region.y ?? 0), Number(region.r ?? region.radius ?? 0), 0, Math.PI * 2);
      ctx.fill();
      ctx.stroke();
    } else if (type === "polygon") {
      this.drawPolygonPath(region.points || []);
      ctx.fill();
      ctx.stroke();
    }
    ctx.restore();
  }

  drawPolygonPath(points) {
    const ctx = this.ctx;
    ctx.beginPath();
    points.forEach((point, index) => {
      const x = Array.isArray(point) ? point[0] : point.x;
      const y = Array.isArray(point) ? point[1] : point.y;
      if (index === 0) ctx.moveTo(x, y);
      else ctx.lineTo(x, y);
    });
    ctx.closePath();
  }

  drawDraft(draft) {
    const region = this.draftToRegion(draft);
    if (!region) return;
    this.drawRegion(region, true);
  }

  drawPolygonDraft() {
    const ctx = this.ctx;
    ctx.save();
    ctx.strokeStyle = "#ffcc00";
    ctx.fillStyle = "#ffcc00";
    ctx.lineWidth = 2;
    ctx.beginPath();
    this.polygonDraft.forEach((point, index) => {
      if (index === 0) ctx.moveTo(point.x, point.y);
      else ctx.lineTo(point.x, point.y);
      ctx.fillRect(point.x - 3, point.y - 3, 6, 6);
    });
    ctx.stroke();
    ctx.restore();
  }
}

app.registerExtension({
  name: "ComfyUI.UniversalRemove.RegionEditor",
  async beforeRegisterNodeDef(nodeType, nodeData) {
    if (nodeData?.name !== "UR Region Editor") return;
    chainCallback(nodeType.prototype, "onNodeCreated", function () {
      this.serialize_widgets = true;
      const editor = new RegionEditorWidget(this);
      this.urRegionEditor = editor;
      const widget = this.addDOMWidget("region_editor", "region_editor", editor.root);
      widget.computeSize = (width) => [width, Math.max(420, this.size?.[1] ? this.size[1] - 160 : 420)];
      this.setSize([Math.max(this.size?.[0] || 680, 680), Math.max(this.size?.[1] || 680, 680)]);
      chainCallback(this, "onResize", function () {
        editor.render();
      });
    });
  },
});