import type { Asset } from "@core/asset";
import type { ImageDocument } from "@core/document";
import type { Rect } from "@core/geometry";
import type { ArtboardId, AssetId, LayerId } from "@core/id";
import type { Layer } from "@core/layer";
import { getLayerMask } from "@core/layer-mask-utils";
import { measureTextLayer } from "@core/text-layer";

export type IndexedLayerInfo = {
  artboardId: ArtboardId;
  parentGroupId?: LayerId;
  layer: Layer;
  siblings: readonly Layer[];
  index: number;
};

export type DocumentReadIndex = {
  assetById: ReadonlyMap<AssetId, Asset>;
  layerById: ReadonlyMap<LayerId, Layer>;
  layerInfoById: ReadonlyMap<LayerId, IndexedLayerInfo>;
  maskLayerIds: ReadonlySet<LayerId>;
  maskLayerIdsByArtboardId: ReadonlyMap<ArtboardId, ReadonlySet<LayerId>>;
  maskLayerIdsByLayerList: ReadonlyMap<readonly Layer[], ReadonlySet<LayerId>>;
  displayLayerCountByArtboardId: ReadonlyMap<ArtboardId, number>;
};

export function createDocumentReadIndex(document: ImageDocument): DocumentReadIndex {
  const assetById = new Map<AssetId, Asset>();
  const layerById = new Map<LayerId, Layer>();
  const layerInfoById = new Map<LayerId, IndexedLayerInfo>();
  const maskLayerIds = new Set<LayerId>();
  const maskLayerIdsByArtboardId = new Map<ArtboardId, ReadonlySet<LayerId>>();
  const maskLayerIdsByLayerList = new Map<readonly Layer[], ReadonlySet<LayerId>>();
  const displayLayerCountByArtboardId = new Map<ArtboardId, number>();

  for (const asset of document.assets) assetById.set(asset.id, asset);

  for (const artboard of document.artboards) {
    const artboardMaskLayerIds = indexLayerTree({
      layers: artboard.layers,
      artboardId: artboard.id,
      layerById,
      layerInfoById,
      documentMaskLayerIds: maskLayerIds,
      maskLayerIdsByLayerList,
    });
    maskLayerIdsByArtboardId.set(artboard.id, artboardMaskLayerIds);
  }

  for (const artboard of document.artboards) {
    displayLayerCountByArtboardId.set(artboard.id, countDisplayLayers(artboard.layers, maskLayerIds));
  }

  return {
    assetById,
    layerById,
    layerInfoById,
    maskLayerIds,
    maskLayerIdsByArtboardId,
    maskLayerIdsByLayerList,
    displayLayerCountByArtboardId,
  };
}

export function forEachLayerBackToFront(layers: readonly Layer[], visit: (layer: Layer) => void) {
  for (let index = layers.length - 1; index >= 0; index -= 1) {
    const layer = layers[index];
    if (layer) visit(layer);
  }
}

export function resolveIndexedLayerBounds(index: DocumentReadIndex, layerOrId: Layer | LayerId): Rect | undefined {
  const layer = typeof layerOrId === "string" ? index.layerById.get(layerOrId) : layerOrId;
  if (!layer) return undefined;

  switch (layer.type) {
    case "group":
      return unionLayerBounds(index, layer.children);
    case "adjustment":
      return undefined;
    case "text": {
      const size = measureTextLayer(layer);
      return { x: layer.transform.position.x, y: layer.transform.position.y, w: size.w * layer.transform.scale.x, h: size.h * layer.transform.scale.y };
    }
    case "image":
    case "raster": {
      const asset = index.assetById.get(layer.assetId);
      if (!asset) return undefined;
      const source = layer.sourceRect ?? { x: 0, y: 0, ...asset.intrinsicSize };

      return {
        x: layer.transform.position.x + source.x * layer.transform.scale.x,
        y: layer.transform.position.y + source.y * layer.transform.scale.y,
        w: source.w * layer.transform.scale.x,
        h: source.h * layer.transform.scale.y,
      };
    }
  }
}

function indexLayerTree(options: {
  layers: readonly Layer[];
  artboardId: ArtboardId;
  parentGroupId?: LayerId;
  layerById: Map<LayerId, Layer>;
  layerInfoById: Map<LayerId, IndexedLayerInfo>;
  documentMaskLayerIds: Set<LayerId>;
  maskLayerIdsByLayerList: Map<readonly Layer[], ReadonlySet<LayerId>>;
}): Set<LayerId> {
  type Frame = { layers: readonly Layer[]; parentGroupId?: LayerId; visited: boolean };
  const stack: Frame[] = [{ layers: options.layers, parentGroupId: options.parentGroupId, visited: false }];
  while (stack.length > 0) {
    const frame = stack.pop();
    if (!frame) continue;
    if (frame.visited) {
      const ids = new Set<LayerId>();
      for (const layer of frame.layers) {
        const mask = getLayerMask(layer);
        if (mask) ids.add(mask.maskLayerId);
        if (layer.type === "group") for (const id of options.maskLayerIdsByLayerList.get(layer.children) ?? []) ids.add(id);
      }
      options.maskLayerIdsByLayerList.set(frame.layers, ids);
      continue;
    }
    stack.push({ ...frame, visited: true });
    for (let index = frame.layers.length - 1; index >= 0; index -= 1) {
      const layer = frame.layers[index];
      if (!layer) continue;
      options.layerById.set(layer.id, layer);
      options.layerInfoById.set(layer.id, { artboardId: options.artboardId, parentGroupId: frame.parentGroupId, layer, siblings: frame.layers, index });
      const mask = getLayerMask(layer);
      if (mask) options.documentMaskLayerIds.add(mask.maskLayerId);
      if (layer.type === "group") stack.push({ layers: layer.children, parentGroupId: layer.id, visited: false });
    }
  }
  return new Set(options.maskLayerIdsByLayerList.get(options.layers) ?? []);
}

function countDisplayLayers(layers: readonly Layer[], maskLayerIds: ReadonlySet<LayerId>): number {
  let count = 0;
  const stack = [...layers];
  while (stack.length > 0) {
    const layer = stack.pop();
    if (!layer) continue;
    if (maskLayerIds.has(layer.id)) continue;
    count += 1;
    if (layer.type === "group") stack.push(...layer.children);
  }
  return count;
}

function unionLayerBounds(index: DocumentReadIndex, layers: readonly Layer[]): Rect | undefined {
  let bounds: Rect | undefined;
  const stack = [...layers];
  while (stack.length > 0) {
    const layer = stack.pop();
    if (!layer) continue;
    if (layer.type === "group") {
      stack.push(...layer.children);
      continue;
    }
    if (layer.type === "adjustment") continue;
    const layerBounds = resolveIndexedLayerBounds(index, layer);
    if (!layerBounds) continue;
    bounds = bounds ? unionRects(bounds, layerBounds) : layerBounds;
  }

  return bounds;
}

function unionRects(a: Rect, b: Rect): Rect {
  const minX = Math.min(a.x, b.x);
  const minY = Math.min(a.y, b.y);
  const maxX = Math.max(a.x + a.w, b.x + b.w);
  const maxY = Math.max(a.y + a.h, b.y + b.h);

  return { x: minX, y: minY, w: maxX - minX, h: maxY - minY };
}