import type { CanvasEdge, CanvasNode, NodeDef, WidgetValue, WorkflowTemplate } from "./types";
import { defaultValues, getDef } from "./registry";

let counter = 0;
export const uid = (prefix = "n") => `${prefix}_${Date.now().toString(36)}${(counter++).toString(36)}`;

export function makeNode(
  id: string,
  type: string,
  x: number,
  y: number,
  values: Record<string, WidgetValue> = {},
  extra: Partial<CanvasNode["data"]> = {},
): CanvasNode {
  const def = getDef(type);
  return {
    id,
    type: "space",
    position: { x, y },
    width: def.width ?? 260,
    data: {
      defType: type,
      values: { ...defaultValues(def), ...values },
      mode: "always",
      status: "idle",
      progress: 0,
      ...extra,
    },
  };
}

export function makeEdge(source: string, sourceHandle: string, target: string, targetHandle: string): CanvasEdge {
  return {
    id: `e_${source}.${sourceHandle}->${target}.${targetHandle}`,
    type: "typed",
    source,
    sourceHandle,
    target,
    targetHandle,
  };
}

/** Edge with its data type filled in from the source def. */
export function typedEdge(nodes: CanvasNode[], source: string, sourceHandle: string, target: string, targetHandle: string): CanvasEdge {
  const src = nodes.find((n) => n.id === source);
  const tgt = nodes.find((n) => n.id === target);
  const sPort = src && getDef(src.data.defType).outputs.find((p) => p.id === sourceHandle);
  const tPort = tgt && getDef(tgt.data.defType).inputs.find((p) => p.id === targetHandle);
  const type = sPort?.type === "any" || !sPort ? (tPort?.type ?? "any") : sPort.type;
  return { ...makeEdge(source, sourceHandle, target, targetHandle), data: { type } };
}

export const FLUX = "black-forest-labs/FLUX.1-schnell";
export const ENHANCER = "finegrain/finegrain-image-enhancer";
export const WHISPER = "hf-audio/whisper-large-v3";

const endpoint = (defs: NodeDef[], apiName: string) => defs.find((d) => d.space?.apiName === apiName) ?? defs[0];
const firstInput = (def: NodeDef | undefined, type: string, fallback: string) => def?.inputs.find((p) => p.type === type)?.id ?? fallback;
const firstOutput = (def: NodeDef | undefined, type: string, fallback: string) => def?.outputs.find((p) => p.type === type)?.id ?? fallback;

const PROMPT = "cinematic photo of a lighthouse on a basalt cliff at blue hour, sea mist, 35mm film";

export const TEMPLATES: WorkflowTemplate[] = [
  {
    id: "txt2img",
    name: "Text to image",
    description: "A prompt into FLUX.1 schnell on ZeroGPU.",
    spaces: [FLUX],
    build: (defsOf) => {
      const flux = endpoint(defsOf(FLUX), "/infer");
      if (!flux) return { nodes: [], edges: [] };
      const nodes = [makeNode("prompt", "in:text", 0, 40, { value: PROMPT }, { title: "Prompt" }), makeNode("flux", flux.type, 380, 0)];
      const edges = [typedEdge(nodes, "prompt", "value", "flux", firstInput(flux, "text", "prompt"))];
      return { nodes, edges };
    },
  },
  {
    id: "enhance",
    name: "Generate & enhance",
    description: "FLUX.1 schnell, then Finegrain's upscaler on the result.",
    spaces: [FLUX, ENHANCER],
    build: (defsOf) => {
      const flux = endpoint(defsOf(FLUX), "/infer");
      const enh = endpoint(defsOf(ENHANCER), "/process");
      if (!flux || !enh) return { nodes: [], edges: [] };
      const nodes = [
        makeNode("prompt", "in:text", 0, 40, { value: PROMPT }, { title: "Prompt" }),
        makeNode("flux", flux.type, 380, 0),
        makeNode("enh", enh.type, 760, 0),
      ];
      const edges = [
        typedEdge(nodes, "prompt", "value", "flux", firstInput(flux, "text", "prompt")),
        typedEdge(nodes, "flux", firstOutput(flux, "image", "result"), "enh", firstInput(enh, "image", "input_image")),
      ];
      return { nodes, edges };
    },
  },
  {
    id: "transcribe",
    name: "Transcribe audio",
    description: "An audio file into Whisper large v3.",
    spaces: [WHISPER],
    build: (defsOf) => {
      const w = endpoint(defsOf(WHISPER), "/transcribe");
      if (!w) return { nodes: [], edges: [] };
      const nodes = [makeNode("audio", "in:audio", 0, 40), makeNode("whisper", w.type, 340, 0)];
      const edges = [typedEdge(nodes, "audio", "value", "whisper", firstInput(w, "audio", "inputs"))];
      return { nodes, edges };
    },
  },
];

export const getTemplate = (id: string) => TEMPLATES.find((t) => t.id === id) ?? TEMPLATES[0];