import {
  iframeSrc,
  showEditor,
  showImageEditor,
  imageUrl,
  imagePrompts,
  targetNode,
  fileName,
  embeddings,
  imagePromptsMulti,
  selectedLayer,
  showLoading,
  loadingCaption,
  alertDialog,
} from "./state.js";
import { van } from "./van.js";
import { app } from "./app.js";
import { api } from "./api.js";
import { Container } from "./Container.js";
import { loadNpyTensor } from "./onnx.js";
import "https://code.iconify.design/3/3.1.0/iconify.min.js";
import { drawSegment, getClicks } from "./LayerEditor.js";

/** @type {import( '../../../web/types/litegraph.js').LGraphGroup} */
const recomputeInsideNodesOps = LGraphGroup.prototype.recomputeInsideNodes;
LGraphGroup.prototype.recomputeInsideNodes = function () {
  this._nodes.length = 0;
  var nodes = this.graph._nodes;
  var node_bounding = new Float32Array(4);

  // const r = recomputeInsideNodesOps.apply(this, arguments);
  for (var i = 0; i < nodes.length; ++i) {
    var node = nodes[i];
    node.getBounding(node_bounding);
    if (!LiteGraph.overlapBounding(this._bounding, node_bounding)) {
      if (node.parentId != undefined && node.parentId == this.id) {
        node.parentId = null;
      }
      continue;
    } //out of the visible area
    this._nodes.push(node);
  }

  this.repositionNodes();
};

const nodeSer = LGraphNode.prototype.serialize;
LGraphNode.prototype.serialize = function () {
  const r = nodeSer.apply(this, arguments);
  r.parentId = this.parentId;
  return r;
};

LGraphGroup.prototype.repositionNodes = function () {
  if (!this.isStack) return;

  const pos = [this.pos[0] + 10, this.pos[1] + 80];

  let height = 0;
  let width = this.size[0];

  let sortedNodes = this._nodes.sort((a, b) => a.pos[1] - b.pos[1]);
  // Separate input and output nodes from the rest
  const inputNodes = sortedNodes.filter(
    (node) => node.properties.routeType === "input"
  );
  const outputNodes = sortedNodes.filter(
    (node) => node.properties.routeType === "output"
  );
  const otherNodes = sortedNodes.filter(
    (node) =>
      node.properties.routeType !== "input" &&
      node.properties.routeType !== "output"
  );

  // Concatenate the arrays so that input nodes are first and output nodes are last
  sortedNodes = [...inputNodes, ...otherNodes, ...outputNodes];

  for (var i = 0; i < sortedNodes.length; ++i) {
    /** @type {LGraphNode} */
    var node = sortedNodes[i];
    node.pos[0] = pos[0];
    node.pos[1] = pos[1] + height;

    node.parentId = this.id;

    if (node.type !== "Reroute") node.size[0] = width - 20;

    if (node.flags.collapsed) {
      height += 40;
    } else if (node.type === "Reroute") {
      height += 30;
      node.pos[1] -= 30;
    } else {
      height += 40;
      height += node.size[1];
    }

    // connect this node output
    if (i < sortedNodes.length - 1) {
      node.disconnectOutput(0);
      node.connect(0, sortedNodes[i + 1], 0);
    }
  }

  this.size[1] = height + 60;
};

const collapseOps = LGraphNode.prototype.collapse;
LGraphNode.prototype.collapse = function (force) {
  collapseOps.apply(this, arguments);
  this.computeParentGroupResize();
};

LGraphNode.prototype.computeParentGroupResize = function () {
  if (this.parentId) {
    const parent = this.graph._groups.find((x) => x.id === this.parentId);
    if (parent) {
      parent.recomputeInsideNodes();
    }
  }
};

const getOpts = LGraphCanvas.prototype.getCanvasMenuOptions;
LGraphCanvas.prototype.getCanvasMenuOptions = function () {
  const r = getOpts.apply(this, arguments);
  r.push({
    content: "Add Stack",
    callback: (info, entry, mouse_event) => {
      var canvas = LGraphCanvas.active_canvas;
      var ref_window = canvas.getCanvasWindow();

      var group = new LiteGraph.LGraphGroup();
      group.pos = canvas.convertEventToCanvasOffset(mouse_event);
      group.isStack = true;
      group.title = "Stack";
      canvas.graph.add(group);

      // add two reroute nodes
      var reroute1 = LiteGraph.createNode("Reroute");
      var reroute2 = LiteGraph.createNode("Reroute");
      reroute1.properties.routeType = "input";
      reroute2.properties.routeType = "output";
      reroute1.pos = [group.pos[0] + 10, group.pos[1] + 40];
      reroute2.pos = [group.pos[0] + 10, group.pos[1] + 70];
      canvas.graph.add(reroute1);
      canvas.graph.add(reroute2);
    },
  });
  return r;
};

const ctor = LGraphGroup.prototype._ctor;
LGraphGroup.prototype._ctor = function (title) {
  ctor.apply(this, arguments);
  this.isStack = false;
  this.id = LiteGraph.uuidv4();
};

const serializationOps = LGraphGroup.prototype.serialize;
LGraphGroup.prototype.serialize = function () {
  const r = serializationOps.apply(this, arguments);
  r.id = this.id;
  r.isStack = this.isStack;
  return r;
};

const configureOps = LGraphGroup.prototype.configure;
LGraphGroup.prototype.configure = function (o) {
  configureOps.apply(this, arguments);
  this.id = o.id;
  this.isStack = o.isStack;
};

/**
 * @typedef {import('../../../web/types/litegraph.js').LGraph} LGraph
 * @typedef {import('../../../web/types/litegraph.js').LGraphNode} LGraphNode
 */

/**
 * Adds a menu handler to a node type.
 * @param {Object} nodeType - The type of the node.
 * @param {Function} cb - The callback function to handle the menu.
 */
function addMenuHandler(nodeType, cb) {
  const getOpts = nodeType.prototype.getExtraMenuOptions;
  nodeType.prototype.getExtraMenuOptions = function () {
    const r = getOpts.apply(this, arguments);
    cb.apply(this, arguments);
    return r;
  };
}

function openInAvatechEditor(url, fileName) {
  let editor = document.getElementById("avatech-editor-iframe");
  iframeSrc.val = url;
  showEditor.val = true;

  editor.contentWindow.postMessage(
    {
      key: "key",
      value: fileName,
      method: "store",
      blendshapes: targetNode.val.widgets.find((x) => x.name === "shape_flow")
        .value,
    },
    "*"
  );
}

function updateBlendshapesPrompts(value) {
  targetNode.val.widgets.find((x) => x.name === "shape_flow").value = value;
  targetNode.val.graph.change();
}

function getInputWidgetValue(node, inputIndex, widgetName) {
  /** @type {LGraph} */
  const graph = app.graph;

  const nodeLink = node.inputs[inputIndex].link;
  if (!nodeLink) return;

  const targetLink = graph.links[nodeLink];

  /** @type {LGraphNode} */
  let nodea = graph._nodes_by_id[targetLink.origin_id];

  while (nodea.type == "Reroute") {
    nodea = nodea.getInputNode(0);
  }

  console.log(targetLink, nodea);
  console.log(nodea.getInputNode(0, true));

  /** @type {string} */
  return nodea.widgets.find((x) => x.name === widgetName).value;
}

/**
 *
 * @param {LGraphNode} node
 */
function showMyImageEditor(node) {
  let connectedImageFileName = getInputWidgetValue(node, 0, "image");
  const split = connectedImageFileName.split("/");
  if (split.length > 1) connectedImageFileName = split[1];

  const embeddingFilename = node.widgets.find(
    (x) => x.name === "embedding_id"
  ).value;

  const v = JSON.parse(
    node.widgets.find((x) => x.name === "image_prompts_json").value
  );

  if (!Array.isArray(v)) {
    // this is a multi prompt
    imagePromptsMulti.val = v;
    selectedLayer.val = Object.keys(imagePromptsMulti.val)[0];
    imagePrompts.val = imagePromptsMulti.val[selectedLayer.val];
  } else {
    // this is a single prompt
    selectedLayer.val = "";
    imagePromptsMulti.val = {};
    imagePrompts.val = v;
  }
  showImageEditor.val = true;
  imageUrl.val = api.apiURL(
    `/view?filename=${encodeURIComponent(
      connectedImageFileName
    )}&type=input&subfolder=${split.length > 1 ? split[0] : ""}`
  );
  const embeedingUrl = api.apiURL(
    `/view?filename=${encodeURIComponent(
      `${embeddingFilename}.npy`
    )}&type=output&subfolder=`
  );
  loadNpyTensor(embeedingUrl).then((tensor) => {
    embeddings.val = tensor;
    drawSegment(getClicks());
  });
  targetNode.val = node;
}

/** @typedef {import('../../../web/types/comfy.js').ComfyExtension} ComfyExtension*/
/** @type {ComfyExtension} */
const ext = {
  getCustomWidgets(app) {
    return {
      SAM_PROMPTS(node, inputName, inputData, app) {
        const btn = node.addWidget("button", "Edit prompt", "", () => {
          let connectedImageFileName = getInputWidgetValue(node, 0, "image");
          if (!connectedImageFileName) {
            alertDialog.val = {
              text: "Please connect an image first",
              time: 3000,
            };
            return;
          }

          loadingCaption.val = "Computing image embedding...";
          showLoading.val = true;

          const split = connectedImageFileName.split("/");
          let id = connectedImageFileName;
          if (split.length > 1) id = split[1];

          node.widgets.find((x) => x.name === "embedding_id").value = id;

          const ckpt = node.widgets.find((x) => x.name === "ckpt").value;

          api
            .fetchApi("/sam_model", {
              method: "POST",
              body: JSON.stringify({
                image: connectedImageFileName,
                embedding_id: id,
                ckpt,
              }),
            })
            .then(() => {
              showLoading.val = false;
              showMyImageEditor(node);
            })
            .catch((err) => {
              console.log(err);
              showLoading.val = false;
            });
        });
        btn.serialize = false;

        return {
          widget: btn,
        };
      },
      BLENDSHAPES_CONFIG(node, inputName, inputData, app) {
        const btn = node.addWidget("button", "Edit Shape Flow", "", () => {
          targetNode.val = node;
          // openInAvatechEditor(
          //   "https://editor.avatech.ai?comfyui=true",
          //   fileName.val
          // );
          openInAvatechEditor("http://localhost:3006?comfyui=true", fileName.val);
        });
        btn.serialize = false;

        return {
          widget: btn,
        };
      },
      MESH_GROUP_CONFIG(node, inputName, inputData, app) {
        const btn = node.addWidget("button", "Add Mesh", "", () => {
          node.addInput("BPY_OBJ" + (node.inputs.length + 1), "BPY_OBJ");
          node.graph.change();
        });
        btn.serialize = false;

        return {
          widget: btn,
        };
      },
      MESH_GROUP_DELETE(node, inputName, inputData, app) {
        const btn = node.addWidget("button", "Delete Mesh", "", () => {
          console.log(node.inputs.length);
          node.removeInput(node.inputs.length - 1);
          node.graph.change();
        });
        btn.serialize = false;

        return {
          widget: btn,
        };
      },
    };
  },

  name: "Avatech.Avatar.BlendshapeEditor",

  init(app) {
    const onNodeMoved = app.canvas.onNodeMoved;
    app.canvas.onNodeMoved = function (node) {
      const r = onNodeMoved?.apply(this, arguments);

      app.graph._groups.forEach((x) => {
        x.recomputeInsideNodes();
      });

      node.computeParentGroupResize();
    };
  },

  async setup() {
    const graphCanvas = document.getElementById("graph-canvas");

    window.addEventListener("message", (event) => {
      if (!event.data.flow || Object.entries(event.data.flow).length <= 0)
        return;
      updateBlendshapesPrompts(event.data.flow);
    });

    api.addEventListener("executed", (evt) => {
      if (evt.detail?.output.gltfFilename) {
        const viewer = document.getElementById(
          "avatech-viewer-iframe"
        ).contentWindow;

        const gltfFilename =
          window.location.protocol +
          "//" +
          api.api_host +
          api.api_base +
          `/view?filename=${evt.detail?.output.gltfFilename[0]}`;

        fileName.val = gltfFilename;
        viewer.postMessage(
          JSON.stringify({
            avatarURL: gltfFilename,
            blendshapes: evt.detail?.output.SHAPE_FLOW[0],
          }),
          "*"
        );
      }
    });

    window.addEventListener(
      "keydown",
      (event) => {
        if (event.key === "Escape") {
          event.preventDefault();
          showImageEditor.val = false;
          showEditor.val = false;
        }
      },
      {
        capture: true,
      }
    );

    graphCanvas.addEventListener("keydown", (event) => {
      if (event.key === "b") {
        event.preventDefault();
        const currentGraph = app.graph.list_of_graphcanvas[0];
        if (currentGraph.selected_nodes.length !== 1) {
          Object.values(currentGraph.selected_nodes).forEach((targetNode) => {
            if (targetNode.mode === 4) targetNode.mode = 0;
            else targetNode.mode = 4;
          });
        } else {
          const targetNode = currentGraph.current_node;
          if (targetNode.mode === 4) targetNode.mode = 0;
          else targetNode.mode = 4;
        }
        app.graph.change();
      }

      if (event.key === "v" && !event.ctrlKey && !event.metaKey) {
        event.preventDefault();
        const currentGraph = app.graph.list_of_graphcanvas[0];
        if (currentGraph.selected_nodes.length !== 1) {
          Object.values(currentGraph.selected_nodes).forEach((targetNode) => {
            if (targetNode.flags.collapsed) targetNode.flags.collapsed = false;
            else targetNode.flags.collapsed = true;

            targetNode.computeParentGroupResize();
          });
        } else {
          const targetNode = currentGraph.current_node;
          console.log(currentGraph.selected_nodes);
          if (targetNode.flags.collapsed) targetNode.flags.collapsed = false;
          else targetNode.flags.collapsed = true;

          targetNode.computeParentGroupResize();
        }
        app.graph.change();
      }
    });

    graphCanvas.addEventListener("keydown", (event) => {
      // if enter
      if (event.key === "Enter") {
        event.preventDefault();
        document.getElementById("queue-button").click();
      }
    });

    van.add(document.body, Container());
  },

  async beforeRegisterNodeDef(nodeType, nodeData, app) {
    switch (nodeData.name) {
      case "ExportGLTF":
        addMenuHandler(nodeType, function (_, options) {
          const output = app.nodeOutputs[this.id + ""];
          if (!output || !output.gltfFilename) return;

          const gltfFilename =
            window.location.protocol +
            "//" +
            api.api_host +
            api.api_base +
            `/view?filename=${output.gltfFilename[0]}`;

          options.unshift({
            content: "Save file",
            callback: () => {
              const a = document.createElement("a");
              let url = new URL(gltfFilename);
              url.searchParams.delete("preview");
              a.href = url;
              a.setAttribute(
                "download",
                new URLSearchParams(url.search).get("filename")
              );
              document.body.append(a);
              a.click();
              requestAnimationFrame(() => a.remove());
            },
          });

          options.unshift({
            content: "Open In Avatech Editor (Local)",
            callback: () => {
              openInAvatechEditor(
                "http://localhost:3006?comfyui=true",
                gltfFilename
              );
            },
          });

          options.unshift({
            content: "Open In Avatech Editor",
            callback: () => {
              openInAvatechEditor(
                "https://editor.avatech.ai?comfyui=true",
                gltfFilename
              );
            },
          });
        });
        break;
      case "SAM_Prompt_Image":
        nodeData.input.required.sam = ["SAM_PROMPTS"];
        // nodeData.input.required.upload = ['IMAGEUPLOAD'];
        // nodeData.input.required.prompts_points = ["IMAGEUPLOAD"];
        addMenuHandler(nodeType, function (_, options) {
          options.unshift({
            content: "Open In Points Editor (Local)",
            callback: () => {
              showMyImageEditor(this);
            },
          });
        });
        break;
      case "SAM":
        nodeData.input.required.sam = ["SAM_PROMPTS"];
        // nodeData.input.required.upload = ['IMAGEUPLOAD'];
        // nodeData.input.required.prompts_points = ["IMAGEUPLOAD"];
        addMenuHandler(nodeType, function (_, options) {
          options.unshift({
            content: "Open In Points Editor (Local)",
            callback: () => {
              showMyImageEditor(this);
            },
          });
        });
        break;
      case "CreateShapeFlow":
        nodeData.input.required.blendshape = ["BLENDSHAPES_CONFIG"];
        addMenuHandler(nodeType, function (_, options) {
          options.unshift({
            content: "Open In Blendshapes Editor",
            callback: () => {
              openInAvatechEditor(
                "https://editor.avatech.ai?comfyui=true",
                gltfFilename
              );
            },
          });
        });
        break;
      case "Mesh_JoinMesh":
        nodeData.input.required.obj = ["MESH_GROUP_CONFIG"];
        nodeData.input.required.del_obj = ["MESH_GROUP_DELETE"];
        break;
      default:
        break;
    }
  },
};

app.registerExtension(ext);