import {
  COLUMN_WIDTH_MIN,
  HEADER_HEIGHT,
  LANE_GAP,
  LANE_MIN_HEIGHT,
  NODE_GAP,
  NODE_HEIGHT,
  NODE_WIDTH,
  NODE_WIDTH_BRANCH,
} from "./traceFlowConstants";
import { TraceGraphLayoutNode, sortNodesInPhase } from "./traceGraphLayout";
import { isFlowGroupContainer } from "./traceDoGroupLayout";

export const CELL_PAD_X = 10;
export const CELL_PAD_Y = 10;

export type CellPlacement = { x: number; y: number; compact: boolean };

const horizWidth = (count: number, nodeW: number) =>
  count * nodeW + Math.max(0, count - 1) * NODE_GAP;

const defaultInnerCellWidth = () => COLUMN_WIDTH_MIN - CELL_PAD_X * 2;

const nodeLayoutSize = (
  node: TraceGraphLayoutNode,
  compact = false,
): { width: number; height: number } => {
  if (isFlowGroupContainer(node)) {
    return {
      width: node.groupBoxWidth ?? NODE_WIDTH,
      height: node.groupBoxHeight ?? NODE_HEIGHT,
    };
  }
  return {
    width: compact ? NODE_WIDTH_BRANCH : NODE_WIDTH,
    height: NODE_HEIGHT,
  };
};

const stackedCellHeight = (members: TraceGraphLayoutNode[]): number =>
  members.reduce((sum, node, index) => {
    const { height } = nodeLayoutSize(node);
    return sum + height + (index > 0 ? NODE_GAP : 0);
  }, 0);

/** Group boxes have variable height — never tile them horizontally with other cards. */
const requiresStackedLayout = (members: TraceGraphLayoutNode[]): boolean =>
  members.length > 1 && members.some(isFlowGroupContainer);

const layoutStackedMembers = (
  sorted: TraceGraphLayoutNode[],
  innerW: number,
  compact: boolean,
): Map<string, CellPlacement> => {
  const result = new Map<string, CellPlacement>();
  let y = CELL_PAD_Y;
  sorted.forEach((node) => {
    const { width, height } = nodeLayoutSize(node, compact);
    result.set(node.id, {
      x: Math.max(0, (innerW - width) / 2),
      y,
      compact,
    });
    y += height + NODE_GAP;
  });
  return result;
};

/** Minimum phase-slot width for a cell's top-level members (excludes DO group children). */
export const cellWidthRequirement = (members: TraceGraphLayoutNode[]): number => {
  const sorted = [...members].sort(sortNodesInPhase);
  const n = sorted.length;
  if (n === 0) return COLUMN_WIDTH_MIN;

  if (n === 1 && isFlowGroupContainer(sorted[0]!)) {
    return Math.max(COLUMN_WIDTH_MIN, (sorted[0]!.groupBoxWidth ?? NODE_WIDTH) + CELL_PAD_X * 2);
  }

  if (n === 1) return COLUMN_WIDTH_MIN;

  if (requiresStackedLayout(sorted)) {
    const maxW = Math.max(...sorted.map((node) => nodeLayoutSize(node).width));
    return Math.max(COLUMN_WIDTH_MIN, maxW + CELL_PAD_X * 2);
  }

  const innerMin = COLUMN_WIDTH_MIN - CELL_PAD_X * 2;
  const branchInner = horizWidth(n, NODE_WIDTH_BRANCH);
  const fullInner = horizWidth(n, NODE_WIDTH);
  const contentInner = Math.max(branchInner, fullInner, innerMin);
  return contentInner + CELL_PAD_X * 2;
};

/** Lay out nodes sharing the same date column + phase lane (side-by-side or vertical stack). */
export const layoutCellMembers = (
  members: TraceGraphLayoutNode[],
  innerCellWidth = defaultInnerCellWidth(),
): Map<string, CellPlacement> => {
  const sorted = [...members].sort(sortNodesInPhase);
  const result = new Map<string, CellPlacement>();
  const n = sorted.length;
  if (n === 0) return result;

  const innerW = innerCellWidth;

  if (n === 1 && isFlowGroupContainer(sorted[0]!)) {
    const w = sorted[0]!.groupBoxWidth ?? NODE_WIDTH;
    result.set(sorted[0]!.id, {
      x: Math.max(0, (innerW - w) / 2),
      y: CELL_PAD_Y,
      compact: false,
    });
    return result;
  }

  if (n === 1) {
    result.set(sorted[0].id, {
      x: (innerW - NODE_WIDTH) / 2,
      y: CELL_PAD_Y,
      compact: false,
    });
    return result;
  }

  if (requiresStackedLayout(sorted)) {
    return layoutStackedMembers(sorted, innerW, false);
  }

  let horizontal = false;
  let nodeW = NODE_WIDTH_BRANCH;
  if (horizWidth(n, NODE_WIDTH_BRANCH) <= innerW) {
    horizontal = true;
    nodeW = NODE_WIDTH_BRANCH;
  } else if (horizWidth(n, NODE_WIDTH) <= innerW) {
    horizontal = true;
    nodeW = NODE_WIDTH;
  }

  if (horizontal) {
    const totalW = horizWidth(n, nodeW);
    const startX = (innerW - totalW) / 2;
    sorted.forEach((node, i) => {
      result.set(node.id, {
        x: startX + i * (nodeW + NODE_GAP),
        y: CELL_PAD_Y,
        compact: true,
      });
    });
    return result;
  }

  return layoutStackedMembers(sorted, innerW, true);
};

export const cellContentHeight = (members: TraceGraphLayoutNode[]): number => {
  const sorted = [...members].sort(sortNodesInPhase);
  const memberCount = sorted.length;
  if (memberCount <= 0) return LANE_MIN_HEIGHT;
  if (memberCount === 1) {
    const sole = sorted[0]!;
    if (isFlowGroupContainer(sole) && sole.groupBoxHeight) {
      return CELL_PAD_Y * 2 + sole.groupBoxHeight;
    }
    return CELL_PAD_Y * 2 + NODE_HEIGHT;
  }

  if (requiresStackedLayout(sorted)) {
    return CELL_PAD_Y * 2 + stackedCellHeight(sorted);
  }

  const innerW = defaultInnerCellWidth();
  const fitsHoriz =
    horizWidth(memberCount, NODE_WIDTH_BRANCH) <= innerW ||
    horizWidth(memberCount, NODE_WIDTH) <= innerW;

  if (fitsHoriz) return CELL_PAD_Y * 2 + NODE_HEIGHT;
  return CELL_PAD_Y * 2 + stackedCellHeight(sorted);
};

export const computeLaneTops = (
  laneCount: number,
  cells: Map<string, TraceGraphLayoutNode[]>,
): { laneTops: number[]; laneHeights: number[]; totalHeight: number } => {
  const laneHeights = Array.from({ length: laneCount }, () => LANE_MIN_HEIGHT);

  cells.forEach((members, key) => {
    const laneIndex = Number(key.split(":")[0]);
    const cellH = cellContentHeight(members);
    laneHeights[laneIndex] = Math.max(laneHeights[laneIndex] ?? LANE_MIN_HEIGHT, cellH);
  });

  const laneTops: number[] = [];
  let y = HEADER_HEIGHT;
  for (let i = 0; i < laneCount; i++) {
    laneTops[i] = y;
    y += laneHeights[i] ?? LANE_MIN_HEIGHT;
    if (i < laneCount - 1) y += LANE_GAP;
  }

  return { laneTops, laneHeights, totalHeight: y + 16 };
};
