paint-ui/store/layer-store.ts

import { create } from 'zustand';
import { v4 as uuidv4 } from 'uuid';
import type { Layer, LayerUpdate, CreateLayerParams, BlendMode } from '@/types';

interface LayerStore {
  /** All layers in the canvas */
  layers: Layer[];
  /** ID of the currently active layer */
  activeLayerId: string | null;

  /** Create a new layer */
  createLayer: (params: CreateLayerParams) => Layer;
  /** Delete a layer by ID */
  deleteLayer: (id: string) => void;
  /** Update layer properties */
  updateLayer: (id: string, updates: LayerUpdate) => void;
  /** Set active layer */
  setActiveLayer: (id: string) => void;
  /** Duplicate a layer */
  duplicateLayer: (id: string) => Layer | null;
  /** Reorder layers */
  reorderLayer: (id: string, newOrder: number) => void;
  /** Merge layer with layer below */
  mergeDown: (id: string) => void;
  /** Flatten all visible layers */
  flattenLayers: () => Layer | null;
  /** Get layer by ID */
  getLayer: (id: string) => Layer | undefined;
  /** Get active layer */
  getActiveLayer: () => Layer | undefined;
  /** Clear all layers */
  clearLayers: () => void;
  /** Add mask to layer */
  addMask: (id: string) => void;
  /** Remove mask from layer */
  removeMask: (id: string) => void;
  /** Enable/disable mask */
  toggleMask: (id: string) => void;
  /** Invert mask */
  invertMask: (id: string) => void;
  /** Apply mask (merge and remove) */
  applyMask: (id: string) => void;
  /** Create a layer group */
  createGroup: (name: string) => Layer;
  /** Add layer to group */
  addToGroup: (layerId: string, groupId: string) => void;
  /** Remove layer from group */
  removeFromGroup: (layerId: string) => void;
  /** Toggle group collapsed state */
  toggleGroupCollapsed: (groupId: string) => void;
  /** Get layers in a group */
  getGroupLayers: (groupId: string) => Layer[];
}

export const useLayerStore = create<LayerStore>((set, get) => ({
  layers: [],
  activeLayerId: null,

  createLayer: (params) => {
    const now = Date.now();
    const layer: Layer = {
      id: uuidv4(),
      name: params.name || `Layer ${get().layers.length + 1}`,
      canvas: null,
      visible: true,
      opacity: params.opacity ?? 1,
      blendMode: params.blendMode || 'normal',
      order: get().layers.length,
      locked: false,
      width: params.width,
      height: params.height,
      x: params.x ?? 0,
      y: params.y ?? 0,
      mask: null,
      groupId: null,
      isGroup: false,
      collapsed: false,
      createdAt: now,
      updatedAt: now,
    };

    // Create canvas element
    const canvas = document.createElement('canvas');
    canvas.width = layer.width;
    canvas.height = layer.height;

    // Fill with color if provided
    if (params.fillColor) {
      const ctx = canvas.getContext('2d');
      if (ctx) {
        ctx.fillStyle = params.fillColor;
        ctx.fillRect(0, 0, layer.width, layer.height);
      }
    }

    layer.canvas = canvas;

    set((state) => ({
      layers: [...state.layers, layer],
      activeLayerId: layer.id,
    }));

    return layer;
  },

  deleteLayer: (id) => {
    set((state) => {
      const newLayers = state.layers.filter((l) => l.id !== id);
      const newActiveId =
        state.activeLayerId === id
          ? newLayers[newLayers.length - 1]?.id || null
          : state.activeLayerId;

      return {
        layers: newLayers,
        activeLayerId: newActiveId,
      };
    });
  },

  updateLayer: (id, updates) => {
    set((state) => ({
      layers: state.layers.map((layer) =>
        layer.id === id
          ? { ...layer, ...updates, updatedAt: Date.now() }
          : layer
      ),
    }));
  },

  setActiveLayer: (id) => {
    set({ activeLayerId: id });
  },

  duplicateLayer: (id) => {
    const layer = get().getLayer(id);
    if (!layer) return null;

    const now = Date.now();
    const newLayer: Layer = {
      ...layer,
      id: uuidv4(),
      name: `${layer.name} copy`,
      order: get().layers.length,
      createdAt: now,
      updatedAt: now,
    };

    // Clone canvas
    if (layer.canvas) {
      const canvas = document.createElement('canvas');
      canvas.width = layer.width;
      canvas.height = layer.height;
      const ctx = canvas.getContext('2d');
      if (ctx) {
        ctx.drawImage(layer.canvas, 0, 0);
      }
      newLayer.canvas = canvas;
    }

    // Clone mask if exists
    if (layer.mask?.canvas) {
      const maskCanvas = document.createElement('canvas');
      maskCanvas.width = layer.width;
      maskCanvas.height = layer.height;
      const maskCtx = maskCanvas.getContext('2d');
      if (maskCtx) {
        maskCtx.drawImage(layer.mask.canvas, 0, 0);
      }
      newLayer.mask = {
        canvas: maskCanvas,
        enabled: layer.mask.enabled,
        inverted: layer.mask.inverted,
      };
    }

    set((state) => ({
      layers: [...state.layers, newLayer],
      activeLayerId: newLayer.id,
    }));

    return newLayer;
  },

  reorderLayer: (id, newOrder) => {
    set((state) => {
      const layers = [...state.layers];
      const layerIndex = layers.findIndex((l) => l.id === id);
      if (layerIndex === -1) return state;

      const [layer] = layers.splice(layerIndex, 1);
      layers.splice(newOrder, 0, layer);

      // Update order values
      return {
        layers: layers.map((l, index) => ({ ...l, order: index })),
      };
    });
  },

  mergeDown: (id) => {
    const layers = get().layers;
    const layerIndex = layers.findIndex((l) => l.id === id);
    if (layerIndex === -1 || layerIndex === 0) return;

    const topLayer = layers[layerIndex];
    const bottomLayer = layers[layerIndex - 1];

    if (!topLayer.canvas || !bottomLayer.canvas) return;

    // Merge onto bottom layer
    const ctx = bottomLayer.canvas.getContext('2d');
    if (!ctx) return;

    ctx.globalAlpha = topLayer.opacity;
    ctx.globalCompositeOperation = topLayer.blendMode as GlobalCompositeOperation;
    ctx.drawImage(topLayer.canvas, topLayer.x, topLayer.y);
    ctx.globalAlpha = 1;
    ctx.globalCompositeOperation = 'source-over';

    // Delete top layer
    get().deleteLayer(id);
  },

  flattenLayers: () => {
    const layers = get().layers.filter((l) => l.visible);
    if (layers.length === 0) return null;

    // Get canvas dimensions
    const width = Math.max(...layers.map((l) => l.width));
    const height = Math.max(...layers.map((l) => l.height));

    // Create flattened layer
    const flatLayer = get().createLayer({
      name: 'Flattened',
      width,
      height,
    });

    if (!flatLayer.canvas) return null;

    const ctx = flatLayer.canvas.getContext('2d');
    if (!ctx) return null;

    // Composite all visible layers
    layers.forEach((layer) => {
      if (layer.canvas) {
        ctx.globalAlpha = layer.opacity;
        ctx.globalCompositeOperation = layer.blendMode as GlobalCompositeOperation;
        ctx.drawImage(layer.canvas, layer.x, layer.y);
      }
    });

    ctx.globalAlpha = 1;
    ctx.globalCompositeOperation = 'source-over';

    // Delete old layers
    layers.forEach((layer) => {
      if (layer.id !== flatLayer.id) {
        get().deleteLayer(layer.id);
      }
    });

    return flatLayer;
  },

  getLayer: (id) => {
    return get().layers.find((l) => l.id === id);
  },

  getActiveLayer: () => {
    const { layers, activeLayerId } = get();
    return layers.find((l) => l.id === activeLayerId);
  },

  clearLayers: () => {
    set({ layers: [], activeLayerId: null });
  },

  addMask: (id) => {
    const layer = get().getLayer(id);
    if (!layer || layer.mask) return;

    // Create mask canvas (white = fully visible)
    const maskCanvas = document.createElement('canvas');
    maskCanvas.width = layer.width;
    maskCanvas.height = layer.height;
    const ctx = maskCanvas.getContext('2d');
    if (ctx) {
      ctx.fillStyle = '#ffffff';
      ctx.fillRect(0, 0, maskCanvas.width, maskCanvas.height);
    }

    get().updateLayer(id, {
      mask: {
        canvas: maskCanvas,
        enabled: true,
        inverted: false,
      },
    });
  },

  removeMask: (id) => {
    get().updateLayer(id, { mask: null });
  },

  toggleMask: (id) => {
    const layer = get().getLayer(id);
    if (!layer?.mask) return;

    get().updateLayer(id, {
      mask: {
        ...layer.mask,
        enabled: !layer.mask.enabled,
      },
    });
  },

  invertMask: (id) => {
    const layer = get().getLayer(id);
    if (!layer?.mask?.canvas) return;

    const maskCtx = layer.mask.canvas.getContext('2d');
    if (!maskCtx) return;

    // Invert the mask canvas
    const imageData = maskCtx.getImageData(0, 0, layer.mask.canvas.width, layer.mask.canvas.height);
    for (let i = 0; i < imageData.data.length; i += 4) {
      imageData.data[i] = 255 - imageData.data[i];     // R
      imageData.data[i + 1] = 255 - imageData.data[i + 1]; // G
      imageData.data[i + 2] = 255 - imageData.data[i + 2]; // B
    }
    maskCtx.putImageData(imageData, 0, 0);

    get().updateLayer(id, {
      mask: {
        ...layer.mask,
        inverted: !layer.mask.inverted,
      },
    });
  },

  applyMask: (id) => {
    const layer = get().getLayer(id);
    if (!layer?.canvas || !layer.mask?.canvas) return;

    const ctx = layer.canvas.getContext('2d');
    const maskCtx = layer.mask.canvas.getContext('2d');
    if (!ctx || !maskCtx) return;

    // Get layer and mask data
    const layerData = ctx.getImageData(0, 0, layer.width, layer.height);
    const maskData = maskCtx.getImageData(0, 0, layer.width, layer.height);

    // Apply mask to layer alpha channel
    for (let i = 0; i < layerData.data.length; i += 4) {
      const maskValue = maskData.data[i]; // Use red channel of mask
      layerData.data[i + 3] = (layerData.data[i + 3] * maskValue) / 255;
    }

    ctx.putImageData(layerData, 0, 0);

    // Remove mask
    get().removeMask(id);
  },

  createGroup: (name) => {
    const now = Date.now();
    const group: Layer = {
      id: uuidv4(),
      name: name || `Group ${get().layers.filter(l => l.isGroup).length + 1}`,
      canvas: null,
      visible: true,
      opacity: 1,
      blendMode: 'normal',
      order: get().layers.length,
      locked: false,
      width: 0,
      height: 0,
      x: 0,
      y: 0,
      mask: null,
      groupId: null,
      isGroup: true,
      collapsed: false,
      createdAt: now,
      updatedAt: now,
    };

    set((state) => ({
      layers: [...state.layers, group],
      activeLayerId: group.id,
    }));

    return group;
  },

  addToGroup: (layerId, groupId) => {
    const layer = get().getLayer(layerId);
    const group = get().getLayer(groupId);

    if (!layer || !group || !group.isGroup) return;

    get().updateLayer(layerId, { groupId });
  },

  removeFromGroup: (layerId) => {
    get().updateLayer(layerId, { groupId: null });
  },

  toggleGroupCollapsed: (groupId) => {
    const group = get().getLayer(groupId);
    if (!group || !group.isGroup) return;

    get().updateLayer(groupId, { collapsed: !group.collapsed });
  },

  getGroupLayers: (groupId) => {
    return get().layers.filter(l => l.groupId === groupId);
  },
}));