// The pictures ComfyUI sends after each sampling step: what the run is
// doing now, which node made the picture, which nodes show it, and the
// finished picture once the run has saved it.
//
// No DOM and no ComfyUI imports, so node:test covers it
// (tests/step_preview.test.mjs). js/shared/step_preview.mjs wires it to the
// server's events and draws it.

import { executionIdOf } from "./graph_ids.mjs";
import { linkById } from "./input_preview.mjs";

// Nothing running, nothing to show.
export const IDLE = Object.freeze({
  status: "idle",
  promptId: null,
  maker: null,
  value: 0,
  max: 0,
  url: null,
  // The finished picture (a /view URL) once a node of this run has saved or
  // previewed it, and how well it matches the steps shown (resultRank).
  result: null,
  resultRank: -1,
  // Between a tiled upscale's split and its stitch every step picture is one
  // tile, not the whole picture (`tiling`). `tilePass`: the steps counted
  // now are a tile's; `whole`: `url` shows the whole picture.
  tiling: false,
  tilePass: false,
  whole: false,
});

const sameId = (a, b) => a !== null && a !== undefined && b !== null && b !== undefined && String(a) === String(b);

// One server event at a time, oldest first; returns the next state and
// never changes the one passed in. Events:
//   { type: "start", promptId }                    a run began
//   { type: "picture", promptId, maker, url, tile } a step picture arrived;
//                                                  `tile`: its maker draws tiles
//   { type: "tiles", promptId, on }                a tiled upscale split (on) or
//                                                  stitched (off) the picture
//   { type: "progress", promptId, node, value, max }
//   { type: "result", promptId, url, rank }        a node of the run showed a picture
//   { type: "end", promptId, kind }                "success" | "interrupted" | "error"
// A step's picture arrives just before its progress report, so a picture
// from a new node starts that node's count at nothing until the report.
// A tile's picture is counted but never replaces a whole one: a tile alone
// looks like a zoomed-in crop of the picture.
export function nextRunState(state, event) {
  const run = state ?? IDLE;
  switch (event?.type) {
    case "start":
      return { ...IDLE, status: "running", promptId: event.promptId ?? null };
    case "picture": {
      if (!event.url) return run;
      // A late picture from a run that already ended changes nothing.
      if (run.status !== "running" && run.status !== "idle" && sameId(event.promptId, run.promptId)) return run;
      const otherRun = run.status !== "running" || (event.promptId && run.promptId && !sameId(event.promptId, run.promptId));
      const base = otherRun ? { ...IDLE, status: "running", promptId: event.promptId ?? null } : run;
      const sameMaker = !otherRun && sameId(base.maker, event.maker);
      const tile = Boolean(base.tiling || event.tile);
      const keep = tile && Boolean(base.url) && base.whole;
      return {
        ...base,
        maker: event.maker ?? null,
        url: keep ? base.url : event.url,
        whole: keep || !tile,
        tilePass: tile,
        value: sameMaker ? base.value : 0,
        max: sameMaker ? base.max : 0,
        // A new step means the picture is being made again (an upscale pass).
        result: null,
        resultRank: -1,
      };
    }
    case "result": {
      if (run.status !== "running" || !event.url) return run;
      if (event.promptId && run.promptId && !sameId(event.promptId, run.promptId)) return run;
      const rank = Number.isFinite(event.rank) ? event.rank : 0;
      if (rank < run.resultRank) return run;
      return { ...run, result: event.url, resultRank: rank };
    }
    case "tiles": {
      if (run.status !== "running") return run;
      if (event.promptId && run.promptId && !sameId(event.promptId, run.promptId)) return run;
      const tiling = Boolean(event.on);
      if (tiling === run.tiling) return run;
      // The upscale starts before its first tile: say so from the split on.
      return tiling ? { ...run, tiling, tilePass: true, value: 0, max: 0 } : { ...run, tiling };
    }
    case "progress": {
      if (run.status !== "running" || !sameId(event.node, run.maker)) return run;
      if (event.promptId && run.promptId && !sameId(event.promptId, run.promptId)) return run;
      const value = Number(event.value);
      const max = Number(event.max);
      if (!Number.isFinite(value) || !Number.isFinite(max)) return run;
      return { ...run, value, max };
    }
    case "end": {
      if (run.status !== "running") return run;
      if (event.promptId && run.promptId && !sameId(event.promptId, run.promptId)) return run;
      const status = event.kind === "interrupted" ? "stopped" : event.kind === "error" ? "failed" : "done";
      return { ...run, status };
    }
    default:
      return run;
  }
}

// "step 3 of 8", or "starting" before the first report; a tiled upscale
// says so ahead of its tile's count.
const TILES_TEXT = "upscaling in tiles";
export function stepText(run) {
  if (!(run?.max > 0)) return run?.tilePass ? TILES_TEXT : "starting";
  const steps = `step ${Math.min(Math.max(0, run.value), run.max)} of ${run.max}`;
  return run.tilePass ? `${TILES_TEXT} · ${steps}` : steps;
}

// What a stopped run leaves on the picture. `saves` adds that the save
// node never got to save it.
export function stoppedText(run, { saves = false } = {}) {
  const where = run?.max > 0 ? `stopped at step ${Math.min(run.value, run.max)} of ${run.max}` : "stopped";
  return saves ? `${where} · nothing saved` : where;
}

// How far along the picture's node is, 0..1.
export function stepFraction(run) {
  if (!(run?.max > 0)) return 0;
  return Math.min(1, Math.max(0, run.value / run.max));
}

// The nodes whose steps are tiles: a tiled upscale's split hands each tile
// to the sampler after it (`TILE_SPLIT` .. `TILE_STITCH`), and Ultimate SD
// Upscale samples its tiles inside one node.
export const TILE_SPLIT = "AUSBOSS_NODES_TiledUpscale";
export const TILE_STITCH = "AUSBOSS_NODES_TiledUpscaleStitch";
const TILE_MAKERS = /^UltimateSDUpscale/;
export function drawsTiles(nodeType) {
  return TILE_MAKERS.test(String(nodeType ?? ""));
}

// How well a picture a node of the run showed stands for the finished
// result: a save or preview node fed by the sampler whose steps were shown
// is best, any other save or preview node next, anything else last.
const RESULT_NODES = /SaveImage|PreviewImage/;
export function resultRank(nodeType, fedByMaker) {
  if (!RESULT_NODES.test(String(nodeType ?? ""))) return 0;
  return fedByMaker ? 2 : 1;
}

// The ids of every node a node's inputs reach back to, on its own graph.
// `outside` says the walk met a link from outside that graph (a box's own
// inputs), which it cannot follow.
export function upstreamIds(graph, start) {
  const ids = new Set();
  let outside = false;
  const queue = [start];
  const seen = new Set([start]);
  while (queue.length) {
    const node = queue.pop();
    for (const input of node?.inputs ?? []) {
      const link = linkById(graph, input?.link);
      if (!link) continue;
      const id = String(link.origin_id);
      ids.add(id);
      const origin = graph?.getNodeById?.(link.origin_id) ?? null;
      if (!origin) {
        outside = true;
        continue;
      }
      if (seen.has(origin)) continue;
      seen.add(origin);
      queue.push(origin);
    }
  }
  return { ids, outside };
}

// Whether the node with this execution id feeds `consumer` (a Save Image),
// so the steps it draws belong in the consumer's picture. A maker in a box
// counts as the box. When the answer cannot be read off the graph (the
// consumer sits in a box the maker is outside of) the steps are shown.
export function feedsInto(rootGraph, consumer, executionId) {
  if (!consumer || executionId === null || executionId === undefined) return false;
  const own = executionIdOf(rootGraph, consumer);
  if (!own) return true;
  const prefix = own.split(":").slice(0, -1);
  const maker = String(executionId).split(":");
  for (let index = 0; index < prefix.length; index += 1) {
    if (maker[index] !== prefix[index]) return true;
  }
  const local = maker[prefix.length];
  if (local === undefined) return true;
  return upstreamIds(consumer.graph, consumer).ids.has(local);
}

// Which of the consumers show a maker's steps: the ones it feeds, or every
// one when it feeds none that the graph can tell (a pack's wireless
// set/get nodes break the trail).
export function consumersOf(rootGraph, consumers, executionId) {
  const list = [...(consumers ?? [])];
  const fed = list.filter((consumer) => feedsInto(rootGraph, consumer, executionId));
  return fed.length ? fed : list;
}