// Models & Pipeline: every node in the live pipeline by group, its values editable, and its shape: swap, add, remove, rewire, regroup.
// The server refuses what would not build (a type that does not fit, a loop) and says why; the reason is shown and nothing changes.
import { composer as c } from "../../composer/composer.js";
import { pickNode } from "./nodepick.js";
import { labelOf, showing, shownValues, sourcesFor } from "./wiring.js";

const SEP = "\0", NEW_GROUP = "\0new";
const fed = (v) => Array.isArray(v);

export const widgetControl = (w, current, set) => {
  const label = w.name;
  if (w.type === "COMBO") {
    const choices = [...(w.choices || [])], unset = current == null || current === "";
    if (!unset && !choices.includes(current)) choices.unshift(current);       // a file gone from disk still shows, so saving something else cannot swap it out
    // Nothing chosen shows as nothing chosen: showing the first file would look picked while the run has none
    // (and with one file there is nothing to change to, so it could never be picked).
    const options = [...(unset ? [{ value: "", label: "— choose —" }] : []), ...choices.map((v) => ({ value: v, label: String(v) }))];
    return choices.length ? c.select.md({ label, value: unset ? "" : current, options, onChange: (v) => { if (v !== "") set(choices.find((x) => String(x) === String(v)) ?? v); } })       // the choice itself: a select gives back text, and ComfyUI refuses "8" for 8
      : c.hint.default({ text: "No choices available — nothing found in the models folder." });
  }
  if (w.type === "INT" || w.type === "FLOAT") {
    return c.number.md({ label, value: current ?? w.default ?? 0, min: w.min, max: w.max, step: w.step ?? (w.type === "INT" ? 1 : undefined),
      precision: w.type === "INT" ? 0 : undefined, onChange: set });
  }
  return w.multiline ? c.textarea.md({ label, value: current ?? "", rows: 4, onCommit: set }) : c.input.md({ label, value: current ?? "", onCommit: set });
};

const FILE = /\.(safetensors|sft|gguf|ckpt|pt|pth|bin|onnx)$/i;
/** Whether a node asks something of the person: a model file to pick, the sampler, or a problem. The rest is plumbing
 *  features need wired in, shown only under "Show all nodes" (v4 kept it inside a fixed graph). */
export const needsChoice = (slot, spec, broken) => broken || !spec || /Loader$/.test(slot.node)
  || (spec.widgets || []).some((w) => w.name === "steps" || w.name === "sampler_name"
    || (w.type === "COMBO" && [...(w.choices || []), (slot.inputs || {})[w.name]].some((v) => typeof v === "string" && FILE.test(v))));
const ALL_KEY = "funui_models_all";

export const models = (app) => function mount() {
  const ps = app.pipeline, api = app.api;
  const page = c.region.stack({ gap: "md" });
  let specs = {}, note = "", group = "", off = null, opened = null, gone = false, openId = null;
  const slotsNow = () => ps.slots() || [];
  const say = (text) => { note = text; draw(); };
  const label = (slot) => labelOf(slot, specs, slotsNow());

  async function describe() {
    const missing = [...new Set(slotsNow().map((s) => s.node))].filter((n) => !(n in specs));
    if (!missing.length) return;
    try { Object.assign(specs, (await api.describeNodes(missing)).nodes || {}); } catch (err) { note = err.message; }
    missing.forEach((n) => { if (!(n in specs)) specs[n] = null; });
  }
  const setInput = (slot, name) => async (value) => {
    const live = slotsNow().find((x) => x.id === slot.id) || slot;           // the slot as it is now, not as drawn: an answer may have replaced it
    if (live.node !== slot.node) return say(`That field belonged to ${slot.node}, which was swapped for ${live.node}: the value was not kept.`);
    if (JSON.stringify((live.inputs || {})[name]) === JSON.stringify(value)) return;          // a blur that changed nothing saves nothing
    await ps.save({ inputs: { [slot.id]: { [name]: value } } });
    const spec = specs[slot.node];
    if (spec && (spec.widgets || []).some((w) => (w.shows || []).some((c) => c.input === name))) {      // a choice that brings fields: what they show is what runs
      const now = slotsNow().find((x) => x.id === slot.id);
      const fill = now && shownValues(now, spec);
      if (fill && Object.keys(fill).length) await ps.save({ inputs: { [slot.id]: fill } });
    }
    say([...(ps.refused() || []), ...ps.saveNotes()].join(" "));
  };
  async function structural(body) {
    const res = await ps.edit(body);
    await describe();
    if (!res.refused.length) {
      const target = body.action === "add" ? slotsNow().slice(-1)[0] : body.action === "replace" || body.action === "unwire" ? slotsNow().find((s) => s.id === body.slot) : null;
      const values = target && shownValues(target, specs[target.node], body.action === "unwire" ? body.input : undefined);
      if (values && Object.keys(values).length) await ps.save({ inputs: { [target.id]: values } });
    }
    say(res.refused.length ? `Not changed: ${res.refused.join(" ")}` : (res.notes || []).join(" "));
    return !res.refused.length;
  }
  const confirm = (title, message, tone) => c.modal.dialogue({ title, message, tone, confirmLabel: title.split(" ")[0] }).result;

  // A select over what could feed an input: "a value" (or not connected), or an output of another node.
  const wiring = (slot, input, type, current, valueLabel) => {
    const options = [{ value: "", label: valueLabel }];
    const sources = sourcesFor(slot, type, slotsNow(), specs), now = fed(current) ? `${current[0]}${SEP}${current[1]}` : "";
    if (now && !sources.some((s) => s.value === now)) options.push({ value: now, label: `${current[0]} · output ${current[1]}` });
    sources.forEach((s) => options.push({ value: s.value, label: `Fed by ${s.label}` }));
    return c.select.md({ label: input, value: now, options, onChange: (v) => {
      if (!v) structural({ action: "unwire", slot: slot.id, input });
      else { const [from, idx] = v.split(SEP); structural({ action: "wire", slot: slot.id, input, from_slot: from, from_output: Number(idx) }); }
    } });
  };

  const groups = () => [...new Set(slotsNow().map((s) => s.group || "Other"))];
  const structure = (slot) => {
    const here = slot.group || "Other";
    const bypass = c.toggle.default({ label: "Bypass", hint: "Skip this node for the run: what it feeds reads what fed it. Its settings stay, switch it back on any time.", checked: Boolean(slot.bypassed), onChange: async (v) => { await ps.setBypass(slot.id, v); draw(); } });
    return c.region.stack({ gap: "sm", children: [bypass, c.field.default({ label: "Group", hint: "Which tab this node sits under. Any node can live in any group.", control: c.toolbar.default({
      items: [c.select.md({ label: "Group", value: here, options: [...groups().map((g) => ({ value: g, label: g })), { value: NEW_GROUP, label: "New group…" }], onChange: async (v) => {
        let name = v;
        if (v === NEW_GROUP) { name = ((await c.modal.prompt({ title: "New group", label: "Name", confirmLabel: "Create" }).result) || "").trim(); if (!name) return draw(); }
        await ps.setGroup(slot.id, name); group = name; draw();
      } })],
      trailing: [c.button.sm({ label: "Swap node…", tone: "ghost", title: "Use a different node here. It must still produce what the nodes after it read.", onClick: async () => {
        const cls = await pickNode(api, `Swap ${label(slot)} for…`);
        if (!cls) return;
        const driven = (slot.roles || []).map((r) => r.label || r.input).filter(Boolean);
        if (await confirm("Swap node", `Swap ${label(slot)} for ${cls}? Everything set on this node (files, values, connections into it) is cleared, and the new node starts empty.${driven.length ? ` The main window drives this node's ${driven.join(", ")}: those controls keep working only for inputs the new node also has, by the same name.` : ""}`, "danger")) structural({ action: "replace", slot: slot.id, node: cls });
      } }), c.button.sm({ label: "Remove", tone: "danger", title: "Take this node out; what it fed is rewired to what fed it, when that is unambiguous.", onClick: async () => {
        if (await confirm("Remove node", `Remove ${label(slot)} from the pipeline?`, "danger")) structural({ action: "remove", slot: slot.id });
      } })] }) })] });
  };

  // Which pack gives a missing node (from ComfyUI-Manager's map, when it is installed), and a button to clone it.
  const providers = {};
  let managerHere = true;
  const installRow = (cls) => {
    const row = c.region.stack({ gap: "sm" });
    const show = () => {
      const url = providers[cls];
      row.set(url === undefined ? [c.hint.default({ text: "Looking for the pack that provides it…" })]
        : url ? [c.settingsRow.default({ label: url.replace(/^https?:\/\/(www\.)?github\.com\//, ""), hint: "Provides this node. Installs into custom_nodes; restart ComfyUI afterwards.", control: c.button.sm({ label: "Install", tone: "primary", onClick: async () => {
          try { await api.pack("install", { url }); c.toast.good({ text: "Installed. Restart ComfyUI for the node to appear." }); } catch (err) { c.toast.warn({ text: err.message }); }
        } }) })]
        : [c.hint.default({ text: managerHere ? "No known pack provides it. Add one by its git URL in Settings ▸ Custom Nodes." : "ComfyUI-Manager is not installed here, so there is no list to look it up in. Add the pack by its git URL in Settings ▸ Custom Nodes." })]);
    };
    show();
    if (!(cls in providers)) api.packProviders([cls]).then((r) => { managerHere = r.manager !== false; providers[cls] = (r.providers || {})[cls] || null; show(); }).catch(() => { providers[cls] = null; show(); });
    return row;
  };

  // The Veda predictors (method veda): the official files that are not in models/veda yet, none at all once both are. The files are not ours to host (MiniMax H3 Community
  // License), so the one way offered is a fetch from the official Hugging Face address on this click, with the licence beside it.
  let vedaSeen = null;          // the last answer: a redrawn page paints it at once, so the row does not empty and fill again
  const vedaRow = () => {
    const row = c.region.stack({ gap: "sm" });
    let timer = null, fetching = false, painted = null;
    const mb = (n) => `${Math.round(n / 1048576)} MB`;
    const paint = (r) => {
      const key = JSON.stringify(r);
      if (key === painted) return;          // nothing changed: leave the rows as they are
      painted = key;
      const busy = r.models.some((m) => m.state === "downloading");
      if (!r.models.some((m) => m.state !== "ready")) return row.set([]);          // both are there: nothing to offer
      row.set([...((r.files || []).length ? [] : [c.banner.info({ text: "Veda needs a predictor file in ComfyUI/models/veda. You can fetch it from Hugging Face; nothing downloads until you click." })]),
        ...r.models.filter((m) => m.state !== "ready").map((m) => c.settingsRow.default({ label: `${m.title} · ${mb(m.size)}`, hint: m.state === "error" ? `Failed: ${m.error}` : m.state === "downloading" ? `${Math.floor(100 * m.done / m.size)}%` : m.hint,
          control: c.button.sm({ label: m.state === "downloading" ? "Downloading…" : m.state === "error" ? "Try again" : "Download", tone: "primary", disabled: busy,
            onClick: async () => { try { await api.vedaDownload(m.id); fetching = true; } catch (err) { c.toast.warn({ text: err.message }); } show(); } }) })),
        c.settingsRow.default({ label: `By downloading you accept the ${r.license.name}.`, hint: "FunUI does not host or bundle the file. It has its own terms of use, including who may use it and where.",
          control: c.button.sm({ label: "Read the licence", tone: "ghost", onClick: () => app.dom.openTab(r.license.url) }) })]);
    };
    const show = async () => {
      let r;
      try { r = await api.vedaState(); } catch (err) { return painted ? undefined : row.set([c.hint.default({ text: `Could not ask about the Veda predictor: ${err.message}` })]); }
      if (gone) return;
      vedaSeen = r;
      const busy = r.models.some((m) => m.state === "downloading");
      if (fetching && !busy) { fetching = false; specs = {}; describe().then(draw); }      // a download ended: the predictor list is read again
      paint(r);
      clearTimeout(timer);
      timer = busy ? setTimeout(() => { if (!gone && row.node && row.node.isConnected) show(); }, 1500) : null;      // a redrawn page leaves its old row behind: that one stops asking
    };
    if (vedaSeen) paint(vedaSeen);
    show();
    return row;
  };

  const slotBlock = (slot) => {
    const spec = specs[slot.node];
    if (spec === undefined) return [c.hint.default({ text: "Loading…" })];
    const head = c.header.sm({ text: label(slot) });
    if (spec === null && (ps.skipped && ps.skipped() || {})[slot.id]) return [head, c.banner.warn({ text: `${slot.node} is skipped because it is not available, so the run goes without it. The Log says why.` })];
    if (spec === null) return [head, structure(slot), c.banner.warn({ text: `${slot.node} is not installed — this slot can't be edited or run. Install the pack that provides it, swap it for another node, or remove it.` }), installRow(slot.node)];
    const sockets = (spec.sockets || []).filter((w) => showing(w, slot, spec)).map((s) => c.field.default({ label: s.name, hint: s.required ? `${s.type} · needed` : s.type, control: wiring(slot, s.name, s.type, slot.inputs && slot.inputs[s.name], "Not connected") }));
    // A field that belongs to a module (the stage a LoRA works in belongs to Two stages) is not drawn while that module is absent or switched off.
    // A field only some models have (SLA is H3's) is not drawn for a model known to lack the trait, nor are the fields that hang off it.
    const have = ps.traits && ps.traits(), gone = new Set((spec.widgets || []).filter((w) => (have && (w.requires || []).some((t) => !have.includes(t))) || (w.needs_module || []).some((id) => ps.hasModule && !ps.hasModule(id))).map((w) => w.name));
    const rows = (spec.widgets || []).filter((w) => !gone.has(w.name) && !(w.shows || []).some((x) => gone.has(x.input)) && showing(w, slot, spec)).flatMap((w) => {
      const cur = slot.inputs && slot.inputs[w.name];
      if (fed(cur)) return [c.field.default({ label: w.name, hint: "Fed by another node: change it there, or pick “A value” to type one here.", control: wiring(slot, w.name, w.type, cur, "A value (typed here)") })];
      const value = cur !== undefined ? cur : w.default;
      const row = w.type === "BOOLEAN" ? c.toggle.default({ label: w.name, hint: w.tooltip || "", checked: Boolean(value), onChange: setInput(slot, w.name) })
        : c.settingsRow.default({ label: w.name, hint: w.tooltip || "", control: widgetControl(w, value, setInput(slot, w.name)) });
      const appOwned = (slot.roles || []).some((r) => r.input === w.name && r.at !== "project.video");       // the app writes it: wiring it elsewhere would make every Generate refuse
      const canTake = !appOwned && (w.type !== "COMBO" || (w.choices || []).length) && sourcesFor(slot, w.type, slotsNow(), specs).length;
      const under = w.name === "veda_predictor" ? [vedaRow()] : [];          // the download goes right under the file it fills
      return canTake ? [row, ...under, c.field.default({ label: "↳ or take it from", hint: "Another node's output can drive this value; one Primitive can drive several inputs.", control: wiring(slot, w.name, w.type, null, "A value (typed here)") })] : [row, ...under];
    });
    return [head, structure(slot), ...sockets, ...rows, ...(!rows.length && !sockets.length ? [c.hint.default({ text: "Nothing to configure on this node." })] : [])];
  };

  // One node as a card: what it is, the one or two values that matter, and whether it needs attention.
  const summary = (slot) => {
    const spec = specs[slot.node];
    if (spec === undefined) return "Loading…";
    if (spec === null) return "Not installed";
    const first = (spec.widgets || [])[0], shown = (spec.widgets || []).map((w) => (slot.inputs || {})[w.name]).filter((v) => typeof v === "string" && v && v !== "default" && v !== "disabled" && !fed(v)).slice(0, 2);
    if (first && first.type === "COMBO" && !(slot.inputs || {})[first.name]) return `${first.name} not set`;
    return shown.length ? shown.join(" · ") : `${(spec.widgets || []).length} setting${(spec.widgets || []).length === 1 ? "" : "s"}`;
  };
  // A server message about this slot: by its id, or refused by the name the person sees. The dot and the banner both use it.
  const mine = (t, id) => t.startsWith(`${id}:`) || t.startsWith(`${id}.`)
    || ((slotsNow().find((s) => s.id === id) || {}).roles || []).some((r) => r.label && t.startsWith(`${r.label} is `));
  const nodeCard = (slot) => {
    const skipped = Boolean((ps.skipped && ps.skipped() || {})[slot.id]);      // optional, and not available: the run goes without it
    const needs = !skipped && (specs[slot.node] === null || ps.incomplete().some((t) => mine(t, slot.id)));
    return c.tile.node({ title: label(slot), sub: skipped ? "skipped: not available" : summary(slot), state: skipped ? "skipped" : needs ? "needs" : "ready",
      onClick: () => { openId = slot.id; draw(); } });
  };
  const addCard = () => c.tile.add({ title: "＋ Add a node", sub: `to ${group}`,
    onClick: async () => { const cls = await pickNode(api, `Add a node to ${group}`); if (cls) structural({ action: "add", node: cls, group }); } });

  // A problem that names no slot first (a loop: "a feeds b feeds a") has no card of its own: its words are shown as they are,
  // and every slot it names counts as broken, so the simple view cannot hide it.
  const unclaimed = () => ps.incomplete().filter((t) => !slotsNow().some((x) => mine(t, x.id)));
  const namedIn = (id, texts) => texts.some((t) => new RegExp(`(^|[^\\w.])${id.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}(?![\\w])`).test(t));
  function status(only) {
    const all = ps.incomplete(), refused = ps.refused(), notes = ps.saveNotes(), loose = unclaimed();
    const incomplete = only ? all.filter((t) => mine(t, only)) : all;
    return [refused.length ? c.banner.warn({ text: `Could not save: ${refused.join(" ")}` }) : null, notes.length ? c.banner.info({ text: notes.join(" ") }) : null,
      incomplete.length ? c.banner.warn({ text: only ? incomplete.join(" ") : `Not ready to generate yet: ${incomplete.length} node${incomplete.length === 1 ? "" : "s"} still need something (the orange dots). Open one to see what.${loose.length ? ` ${loose.join(" ")}` : ""}` }) : ps.queueable() && !only ? c.hint.default({ text: "Every slot is filled — this pipeline is ready to generate." }) : null];
  }

  const tools = () => c.toolbar.default({ items: [c.button.sm({ label: "Start from…", tone: "ghost", title: "Replace this pipeline with a model's own starting point", onClick: presets }),
    c.button.sm({ label: "Import workflow…", tone: "ghost", title: "Use a ComfyUI workflow (UI or API export) as this pipeline", onClick: importWorkflow }),
    c.button.sm({ label: "Revert", tone: "ghost", disabled: !opened, title: "Put back what was here when this page was opened", onClick: async () => { const r = await ps.restore(opened); say(r.refused.length ? r.refused[0] : ""); } })],
    trailing: [c.button.sm({ label: "🖼 Export settings…", tone: "ghost", title: "Render this pipeline as a PNG: loaders, typed-in values, torch / CUDA / attention", onClick: card })] });

  async function presets() {
    let found;
    try { found = (await api.pipelinePresets()).presets || []; } catch (err) { return say(err.message); }
    if (!found.length) return say("No module offers a starting point.");
    const id = await c.modal.choice({ title: "Start from…", items: found.map((p) => ({ id: p.id, label: p.title, hint: p.module })) }).result;
    const pick = found.find((p) => p.id === id);
    if (!pick || !(await confirm("Replace pipeline", `Replace the whole pipeline with “${pick.title}”? Every file and value you set here is cleared. Revert puts it back until you leave this page.`, "danger"))) return;
    const r = await ps.restore({ slots: pick.slots, removed: [], unwired: {} });
    if (!r.refused.length) {      // project settings the preset goes with (render steps, a picture model's one frame)
      Object.entries(pick.prefs || {}).forEach(([k, v]) => app.project.setPref(k, v));
      Object.entries(pick.project || {}).forEach(([k, v]) => app.project.setField(k, v));
    }
    specs = {}; await describe(); say(r.refused.length ? `Not changed: ${r.refused[0]}` : "");
  }
  async function importWorkflow() {
    const file = await app.dom.pickFile(".json,application/json");
    if (!file) return;
    let got;
    try { got = await api.importWorkflow(JSON.parse(await file.text())); } catch (err) { return say(err instanceof SyntaxError ? "That file is not valid JSON." : err.message); }
    const names = { prompt: "Prompt", negative: "Negative prompt", seed: "Seed", width: "Width", height: "Height", frames: "Length", fps: "FPS", image: "Start picture" };
    const lines = [`${got.slots.length} nodes.`, ...Object.entries(names).map(([k, label]) => `${label}: ${got.bound[k] ? "→ " + got.bound[k].join(", ") : "not connected (set on the node itself)"}`), ...got.notes];
    if (!(await confirm("Use this workflow", `${lines.join("\n")}\n\nReplace the whole pipeline with it? Revert puts the old one back until you leave this page.`, "danger"))) return;
    const r = await ps.restore({ slots: got.slots, removed: [], unwired: {} });
    specs = {}; await describe(); say(r.refused.length ? `Not changed: ${r.refused[0]}` : "");
  }
  async function card() {
    let url;
    try { url = URL.createObjectURL(await api.settingsCard(slotsNow(), app.project.project && app.project.project.name, app.theme && app.theme.get ? app.theme.get() : "dark")); } catch (err) { return say(`Could not render the card: ${err.message}`); }
    const win = c.modal.generic({ title: "Pipeline settings", size: "lg", body: c.viewer.media({ src: url, kind: "image" }), onClose: () => URL.revokeObjectURL(url) });
    win.setFooter({ actions: [c.button.sm({ label: "Save PNG", tone: "primary", onClick: () => app.dom.download(url, "funui-settings.png") })] });
  }

  function draw() {
    if (gone) return;
    if (ps.loading()) return page.set([c.hint.default({ text: "Loading…" })]);
    if (ps.loadError()) return page.set([c.banner.warn({ text: `Could not load models & pipeline: ${ps.loadError()}` })]);
    const all = groups(), slots = slotsNow();
    if (!all.includes(group)) group = all[0] || "";
    const open = openId && slots.find((x) => x.id === openId);
    if (openId && !open) openId = null;                       // removed or swapped away
    if (open) {
      const back = c.button.sm({ label: `‹ ${open.group || "Other"}`, tone: "ghost", onClick: () => { openId = null; draw(); } });
      const detail = c.card.default({ children: slotBlock(open) });
      return page.set([back, note ? c.banner.warn({ text: note }) : null, ...status(open.id), detail].filter(Boolean));
    }
    let everything = false;
    try { everything = localStorage.getItem(ALL_KEY) === "1"; } catch { /* the simple view */ }
    const show = (on) => { try { on ? localStorage.setItem(ALL_KEY, "1") : localStorage.removeItem(ALL_KEY); } catch { /* this visit only */ } draw(); };
    const loose = unclaimed();
    const chosen = slots.filter((x) => needsChoice(x, specs[x.node], specs[x.node] === null || ps.incomplete().some((t) => mine(t, x.id)) || namedIn(x.id, loose)));
    const grid = c.grid.cards({ children: everything
      ? [...slots.filter((x) => (x.group || "Other") === group).map(nodeCard), group ? addCard() : null].filter(Boolean)
      : chosen.map(nodeCard) });
    const rest = slots.length - chosen.length;
    page.set([tools(), note ? c.banner.warn({ text: note }) : null, ...status(),
      !slots.length ? c.emptyState.default({ icon: "⬡", title: "No pipeline loaded", hint: "Is ComfyUI reachable? Or start from a model's starting point above." }) : null,
      slots.length ? c.toggle.default({ label: "Show all nodes", hint: everything ? "Every node, by group: add, swap, rewire." : `Only what needs a choice: model files, the sampler, anything broken.${rest ? ` ${rest} more only connect features, nothing to pick there.` : ""}`, checked: everything, onChange: show }) : null,
      everything && all.length ? c.tabs.underline({ label: "Group", tabs: all.map((g) => ({ value: g, label: g })), value: group, onChange: (g) => { group = g; draw(); } }) : null,
      grid].filter(Boolean));
  }
  // A redraw while a box has focus would eat what is being typed: it waits for the focus to leave.
  const typing = () => app.dom.typingIn(page.node);
  page.node.addEventListener("focusout", () => setTimeout(() => { if (!typing()) draw(); }));
  draw();
  ps.ensureLoaded().then(async () => { opened = ps.snapshot(); await describe(); draw(); });
  off = ps.subscribe((_s, kind) => {
    if (kind === "load") opened = ps.snapshot();          // another project's pipeline went in: Revert goes back to that one
    if (!typing()) { describe().then(draw); }
  });
  return { node: page.node, destroy: () => { gone = true; off && off(); page.node.remove(); } };
};