import { commandIds } from "@commands/ids";
import type { ImageDocument } from "@core/document";
import type { Transform } from "@core/geometry";
import type { Layer } from "@core/layer";
import type { AppStore } from "@editor/store";
import type { GenerationCandidate, SelectionState, ViewportState } from "@editor/state";
import type { GenerateSettings } from "@editor/tools";
import { buildInpaintBundle, type InpaintBundle } from "./inpaintPrep";

export async function runGenerate(options: {
  document: ImageDocument;
  selection: SelectionState;
  viewport: ViewportState;
  settings: GenerateSettings;
  dispatch: AppStore["dispatch"];
}) {
  const { document, selection, settings, dispatch } = options;
  const artboard = selection.artboardId ? document.artboards.find((candidate) => candidate.id === selection.artboardId) : document.artboards[0];
  if (!artboard) return;

  const target = resolveSelectedImage(document, selection);
  const inpaintBundle = settings.mode === "inpaint" ? await buildInpaintBundle(document, selection, settings) : undefined;
  const inputImage = inpaintBundle?.inputImage ?? (target && settings.mode !== "text-to-image" ? await imageSourceToDataUrl(target.asset.source) : undefined);
  const maskImage = inpaintBundle?.maskImage;
  const seed = resolveSeed(settings.seed);
  const requestSettings = { ...settings, seed };
  const width = inpaintBundle?.width ?? settings.width;
  const height = inpaintBundle?.height ?? settings.height;
  const generated = await requestGenerate({
    settings: requestSettings,
    width,
    height,
    inputImage,
    maskImage,
    inpaintBundle,
  });
  const intrinsicSize = await loadImageSize(generated.source);
  const targetArtboardId = inpaintBundle?.placement.artboardId ?? artboard.id;
  const placement = {
    artboardId: targetArtboardId,
    layerName: inpaintBundle?.placement.layerName ?? "Generated image",
    transform: generatedLayerTransform(inpaintBundle, intrinsicSize),
  };

  dispatch(commandIds.generationAddCandidate, {
    candidate: createGenerationCandidate({
      source: generated.source,
      mimeType: generated.mimeType,
      intrinsicSize,
      settings: requestSettings,
      seed,
      width,
      height,
      inputImage,
      maskImage,
      placement,
      inpaintBundle,
    }),
  });
}

export async function runGenerateFromCandidate(options: {
  candidate: GenerationCandidate;
  settings?: GenerateSettings;
  dispatch: AppStore["dispatch"];
}) {
  const settings = options.settings ?? options.candidate.settings;
  const seed = resolveSeed(settings.seed);
  const requestSettings = { ...settings, seed };
  const generated = await requestGenerate({
    settings: requestSettings,
    width: options.candidate.width,
    height: options.candidate.height,
    inputImage: options.candidate.inputImage,
    maskImage: options.candidate.maskImage,
    inpaintCandidate: options.candidate,
  });
  const intrinsicSize = await loadImageSize(generated.source);

  options.dispatch(commandIds.generationAddCandidate, {
    candidate: {
      ...options.candidate,
      id: crypto.randomUUID(),
      source: generated.source,
      mimeType: generated.mimeType,
      intrinsicSize,
      settings: requestSettings,
      seed,
    },
  });
}

function createGenerationCandidate(options: {
  source: string;
  mimeType: string;
  intrinsicSize: { w: number; h: number };
  settings: GenerateSettings;
  seed: number;
  width: number;
  height: number;
  inputImage?: string;
  maskImage?: string;
  placement: GenerationCandidate["placement"];
  inpaintBundle?: InpaintBundle;
}): GenerationCandidate {
  return {
    id: crypto.randomUUID(),
    source: options.source,
    mimeType: options.mimeType,
    intrinsicSize: options.intrinsicSize,
    mode: options.settings.mode,
    settings: options.settings,
    seed: options.seed,
    width: options.width,
    height: options.height,
    inputImage: options.inputImage,
    maskImage: options.maskImage,
    placement: options.placement,
    inpaint: options.inpaintBundle
      ? {
          targetLayerId: options.inpaintBundle.targetLayerId,
          maskLayerId: options.inpaintBundle.maskLayerId,
          sourceAssetId: options.inpaintBundle.sourceAssetId,
          maskAssetId: options.inpaintBundle.maskAssetId,
          inputImage: options.inpaintBundle.inputImage,
          maskImage: options.inpaintBundle.maskImage,
          crop: options.inpaintBundle.crop,
          mask: options.inpaintBundle.mask,
          backend: options.inpaintBundle.backend,
        }
      : undefined,
  };
}

async function requestGenerate(options: {
  settings: GenerateSettings;
  width: number;
  height: number;
  inputImage?: string;
  maskImage?: string;
  inpaintBundle?: InpaintBundle;
  inpaintCandidate?: GenerationCandidate;
}) {
  const response = await fetch("/api/comfy/generate", {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({
      mode: options.settings.mode,
      model: options.settings.model,
      prompt: options.settings.prompt,
      negativePrompt: options.settings.negativePrompt,
      strength: options.settings.strength,
      steps: options.settings.steps,
      cfg: options.settings.cfg,
      seed: options.settings.seed,
      sampler: options.settings.sampler,
      scheduler: options.settings.scheduler,
      width: options.width,
      height: options.height,
      outpaint: options.settings.outpaint,
      inpaint: resolveInpaintRequest(options.inpaintBundle, options.inpaintCandidate, options.settings),
      inputImage: options.inputImage,
      maskImage: options.maskImage,
    }),
  });
  if (!response.ok) throw new Error(await response.text());
  return await response.json() as { source: string; mimeType: string };
}

function resolveInpaintRequest(inpaintBundle: InpaintBundle | undefined, inpaintCandidate: GenerationCandidate | undefined, settings: GenerateSettings) {
  if (inpaintBundle) {
    return {
      growMaskBy: inpaintBundle.backend.growMaskBy,
      maskedContent: inpaintBundle.backend.maskedContent,
      maskBlur: inpaintBundle.backend.maskBlur,
      maskFeather: inpaintBundle.backend.maskFeather,
      maskExpand: inpaintBundle.backend.maskExpand,
      cropPadding: inpaintBundle.backend.cropPadding,
      maskPolarity: inpaintBundle.mask.polarity,
      crop: inpaintBundle.crop,
      placement: inpaintBundle.placement,
    };
  }

  if (inpaintCandidate?.inpaint) {
    return {
      growMaskBy: inpaintCandidate.inpaint.backend.growMaskBy,
      maskedContent: inpaintCandidate.inpaint.backend.maskedContent,
      maskBlur: inpaintCandidate.inpaint.backend.maskBlur,
      maskFeather: inpaintCandidate.inpaint.backend.maskFeather,
      maskExpand: inpaintCandidate.inpaint.backend.maskExpand,
      cropPadding: inpaintCandidate.inpaint.backend.cropPadding,
      maskPolarity: inpaintCandidate.inpaint.mask.polarity,
      crop: inpaintCandidate.inpaint.crop,
      placement: inpaintCandidate.placement,
    };
  }

  return settings.inpaint;
}

function generatedLayerTransform(inpaintBundle: InpaintBundle | undefined, intrinsicSize: { w: number; h: number }): Transform {
  if (!inpaintBundle) return { position: { x: 0, y: 0 }, scale: { x: 1, y: 1 }, rotation: 0 };
  return {
    position: { ...inpaintBundle.placement.transform.position },
    scale: {
      x: (inpaintBundle.placement.transform.scale.x * inpaintBundle.width) / Math.max(1, intrinsicSize.w),
      y: (inpaintBundle.placement.transform.scale.y * inpaintBundle.height) / Math.max(1, intrinsicSize.h),
    },
    rotation: inpaintBundle.placement.transform.rotation,
  };
}

function resolveSeed(seed: number): number {
  return seed < 0 ? Math.floor(Math.random() * 2 ** 32) : Math.round(seed);
}

function resolveSelectedImage(document: ImageDocument, selection: SelectionState) {
  const layerId = selection.layerIds[0];
  if (!layerId) return undefined;
  const layer = findLayer(document.artboards.find((artboard) => artboard.id === selection.artboardId)?.layers ?? [], layerId);
  if (!layer || layer.type === "group") return undefined;
  const asset = document.assets.find((candidate) => candidate.id === layer.assetId);
  const maskLayer = layer.clippingMask ? findLayer(document.artboards.flatMap((artboard) => artboard.layers), layer.clippingMask.maskLayerId) : undefined;
  const maskAsset = maskLayer && maskLayer.type !== "group" ? document.assets.find((candidate) => candidate.id === maskLayer.assetId) : undefined;
  return asset ? { layer, asset, maskAsset } : undefined;
}

function findLayer(layers: readonly Layer[], layerId: string): Layer | undefined {
  for (const layer of layers) {
    if (layer.id === layerId) return layer;
    if (layer.type === "group") {
      const found = findLayer(layer.children, layerId);
      if (found) return found;
    }
  }
  return undefined;
}

async function imageSourceToDataUrl(source: string) {
  if (source.startsWith("data:")) return source;
  const image = await loadImage(source);
  const canvas = document.createElement("canvas");
  canvas.width = image.naturalWidth;
  canvas.height = image.naturalHeight;
  const context = canvas.getContext("2d");
  if (!context) throw new Error("Unable to read selected image");
  context.drawImage(image, 0, 0);
  return canvas.toDataURL("image/png");
}

function loadImageSize(source: string): Promise<{ w: number; h: number }> {
  return loadImage(source).then((image) => ({ w: image.naturalWidth, h: image.naturalHeight }));
}

function loadImage(source: string): Promise<HTMLImageElement> {
  return new Promise((resolve, reject) => {
    const image = new Image();
    image.onload = () => resolve(image);
    image.onerror = () => reject(new Error("Failed to load image"));
    image.src = source;
  });
}