import type { Asset } from "@core/asset";
import type { AssetGenerationProvenance, GeneratedAssetAcceptance } from "@core/asset-provenance";
import type { ImageDocument } from "@core/document";
import type { ArtboardId, AssetId, LayerId } from "@core/id";
import type { ImageLayer } from "@core/image-layer";
import type { Layer } from "@core/layer";
import type { GenerationCandidate, GenerationCompareMode, GenerationState } from "@editor/state";
import type { Command } from "./command";
import { commandIds } from "./ids";

export type GenerationAddCandidatePayload = {
  candidate: GenerationCandidate;
};

export type GenerationSelectCandidatePayload = {
  candidateId?: string;
};

export type GenerationSetCompareModePayload = {
  mode: GenerationCompareMode;
};

export type GenerationRemoveCandidatePayload = {
  candidateId: string;
};

export type GenerationApplyCandidateAsLayerPayload = {
  candidateId: string;
  assetId: AssetId;
  layerId: LayerId;
  variant?: boolean;
};

export type GenerationReplaceCandidatePixelsPayload = {
  candidateId: string;
  source: string;
  mimeType?: string;
};

const maxCandidates = 12;
const generationCompareModes = new Set<GenerationCompareMode>(["result", "before", "split"]);

export const generationAddCandidateCommand: Command<GenerationAddCandidatePayload> = {
  id: commandIds.generationAddCandidate,
  name: "Add generation candidate",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    const candidates = [payload.candidate, ...state.editor.generation.candidates.filter((candidate) => candidate.id !== payload.candidate.id)].slice(0, maxCandidates);
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: {
          candidates,
          selectedCandidateId: payload.candidate.id,
          compareMode: "result",
        },
      },
    };
  },
};

export const generationSelectCandidateCommand: Command<GenerationSelectCandidatePayload> = {
  id: commandIds.generationSelectCandidate,
  name: "Select generation candidate",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    const selectedCandidateId = payload.candidateId && state.editor.generation.candidates.some((candidate) => candidate.id === payload.candidateId) ? payload.candidateId : undefined;
    if (state.editor.generation.selectedCandidateId === selectedCandidateId) return state;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: {
          ...state.editor.generation,
          selectedCandidateId,
        },
      },
    };
  },
};

export const generationSetCompareModeCommand: Command<GenerationSetCompareModePayload> = {
  id: commandIds.generationSetCompareMode,
  name: "Set generation compare mode",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    if (!generationCompareModes.has(payload.mode)) return state;
    if (state.editor.generation.compareMode === payload.mode) return state;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: {
          ...state.editor.generation,
          compareMode: payload.mode,
        },
      },
    };
  },
};

export const generationRemoveCandidateCommand: Command<GenerationRemoveCandidatePayload> = {
  id: commandIds.generationRemoveCandidate,
  name: "Remove generation candidate",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    const candidates = state.editor.generation.candidates.filter((candidate) => candidate.id !== payload.candidateId);
    if (candidates.length === state.editor.generation.candidates.length) return state;
    const selectedCandidateId = state.editor.generation.selectedCandidateId === payload.candidateId ? candidates[0]?.id : state.editor.generation.selectedCandidateId;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: {
          candidates,
          selectedCandidateId,
          compareMode: candidates.length > 0 ? state.editor.generation.compareMode : "result",
        },
      },
    };
  },
};

export const generationClearCandidatesCommand: Command = {
  id: commandIds.generationClearCandidates,
  name: "Clear generation candidates",
  history: { mode: "ignore" },
  execute({ state }) {
    if (state.editor.generation.candidates.length === 0 && !state.editor.generation.selectedCandidateId && state.editor.generation.compareMode === "result") return state;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: { candidates: [], selectedCandidateId: undefined, compareMode: "result" },
      },
    };
  },
};

export const generationApplyCandidateAsLayerCommand: Command<GenerationApplyCandidateAsLayerPayload> = {
  id: commandIds.generationApplyCandidateAsLayer,
  name: "Apply generation candidate as layer",
  execute({ state }, payload) {
    const candidate = state.editor.generation.candidates.find((item) => item.id === payload.candidateId);
    if (!candidate) return state;
    if (state.document.assets.some((asset) => asset.id === payload.assetId) || findLayerLocation(state.document, payload.layerId)) return state;
    if (!state.document.artboards.some((artboard) => artboard.id === candidate.placement.artboardId)) return state;

    const asset: Asset = {
      id: payload.assetId,
      name: payload.variant ? `${candidate.placement.layerName} variant` : candidate.placement.layerName,
      mimeType: candidate.mimeType,
      source: candidate.source,
      intrinsicSize: { ...candidate.intrinsicSize },
      provenance: generationProvenance(candidate, payload.variant ? "variant-layer" : "layer"),
    };
    const layer: ImageLayer = {
      id: payload.layerId,
      type: "image",
      name: payload.variant ? `${candidate.placement.layerName} variant` : candidate.placement.layerName,
      visible: true,
      locked: false,
      opacity: 1,
      assetId: asset.id,
      transform: {
        position: { ...candidate.placement.transform.position },
        scale: { ...candidate.placement.transform.scale },
        rotation: candidate.placement.transform.rotation,
      },
    };

    return {
      ...state,
      document: insertLayerAtTop({ ...state.document, assets: [...state.document.assets, asset] }, candidate.placement.artboardId, layer),
      editor: {
        ...state.editor,
        generation: clearCommittedGenerationPreview(state.editor.generation),
        selection: { artboardId: candidate.placement.artboardId, layerIds: [layer.id] },
      },
    };
  },
};

export const generationReplaceCandidatePixelsCommand: Command<GenerationReplaceCandidatePixelsPayload> = {
  id: commandIds.generationReplaceCandidatePixels,
  name: "Replace masked pixels with generation candidate",
  execute({ state }, payload) {
    const candidate = state.editor.generation.candidates.find((item) => item.id === payload.candidateId);
    if (!candidate?.inpaint || !payload.source.trim()) return state;
    const targetAsset = state.document.assets.find((asset) => asset.id === candidate.inpaint?.sourceAssetId);
    const targetLayerLocation = findLayerLocation(state.document, candidate.inpaint.targetLayerId);
    if (!targetAsset || !targetLayerLocation) return state;

    return {
      ...state,
      document: {
        ...state.document,
        assets: state.document.assets.map((asset) =>
          asset.id === targetAsset.id
            ? {
                ...asset,
                source: payload.source,
                mimeType: payload.mimeType ?? asset.mimeType,
                provenance: generationProvenance(candidate, "replacement"),
              }
            : asset,
        ),
      },
      editor: {
        ...state.editor,
        generation: clearCommittedGenerationPreview(state.editor.generation),
        selection: { artboardId: targetLayerLocation.artboardId, layerIds: [candidate.inpaint.targetLayerId] },
      },
    };
  },
};

export const generationCommands = [
  generationAddCandidateCommand,
  generationSelectCandidateCommand,
  generationSetCompareModeCommand,
  generationRemoveCandidateCommand,
  generationClearCandidatesCommand,
  generationApplyCandidateAsLayerCommand,
  generationReplaceCandidatePixelsCommand,
] satisfies Command<unknown>[];

type LayerLocation = {
  artboardId: ArtboardId;
  layer: Layer;
};

function insertLayerAtTop(document: ImageDocument, artboardId: ArtboardId, layer: Layer): ImageDocument {
  return {
    ...document,
    artboards: document.artboards.map((artboard) => artboard.id === artboardId ? { ...artboard, layers: [layer, ...artboard.layers] } : artboard),
  };
}

function findLayerLocation(document: ImageDocument, layerId: LayerId): LayerLocation | undefined {
  for (const artboard of document.artboards) {
    const layer = findLayerInTree(artboard.layers, layerId);
    if (layer) return { artboardId: artboard.id, layer };
  }
  return undefined;
}

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

function clearCommittedGenerationPreview(generation: GenerationState): GenerationState {
  if (generation.candidates.length === 0 && !generation.selectedCandidateId && generation.compareMode === "result") return generation;
  return { candidates: [], selectedCandidateId: undefined, compareMode: "result" };
}

function generationProvenance(candidate: GenerationCandidate, acceptance: GeneratedAssetAcceptance): AssetGenerationProvenance {
  return {
    kind: "generated",
    candidateId: candidate.id,
    mode: candidate.mode,
    acceptance,
    prompt: candidate.settings.prompt,
    negativePrompt: candidate.settings.negativePrompt,
    seed: candidate.seed,
    outputSize: { ...candidate.intrinsicSize },
    settings: {
      architecture: candidate.settings.architecture,
      model: candidate.settings.model,
      textEncoder: candidate.settings.textEncoder,
      vae: candidate.settings.vae,
      strength: candidate.settings.strength,
      steps: candidate.settings.steps,
      cfg: candidate.settings.cfg,
      sampler: candidate.settings.sampler,
      scheduler: candidate.settings.scheduler,
      width: candidate.settings.width,
      height: candidate.settings.height,
    },
    inpaint: candidate.inpaint
      ? {
          targetLayerId: candidate.inpaint.targetLayerId,
          maskLayerId: candidate.inpaint.maskLayerId,
          sourceAssetId: candidate.inpaint.sourceAssetId,
          maskAssetId: candidate.inpaint.maskAssetId,
          crop: {
            assetBounds: { ...candidate.inpaint.crop.assetBounds },
            documentBounds: { ...candidate.inpaint.crop.documentBounds },
            padding: candidate.inpaint.crop.padding,
            maskedAreaOnly: candidate.inpaint.crop.maskedAreaOnly,
          },
          mask: {
            polarity: candidate.inpaint.mask.polarity,
            activeBounds: { ...candidate.inpaint.mask.activeBounds },
          },
          backend: {
            growMaskBy: candidate.inpaint.backend.growMaskBy,
            maskedContent: candidate.inpaint.backend.maskedContent,
            maskBlur: candidate.inpaint.backend.maskBlur,
            maskFeather: candidate.inpaint.backend.maskFeather,
            maskExpand: candidate.inpaint.backend.maskExpand,
            cropPadding: candidate.inpaint.backend.cropPadding,
          },
        }
      : undefined,
  };
}