import type { Node, Edge } from "@xyflow/react";
import { SEQ_SIZES, SEQ_LAYOUT } from "../types/sequence/constants";

/**
 * Custom layout for sequence diagrams.
 * Positions participants horizontally, computes message Y from edge order,
 * sizes participant nodes to full lifeline height, computes activation bars,
 * and positions fragments around their message ranges.
 *
 * Returns new node/edge objects — does NOT mutate inputs.
 *
 * NOTE: This is synchronous — sequence diagrams are small (no Web Worker needed).
 */
export function computeSequenceLayout(
  nodes: Node[],
  edges: Edge[],
): { nodes: Node[]; edges: Edge[] } {
  const participants = nodes.filter((n) => n.type === "seqParticipant");
  const fragments = nodes.filter((n) => n.type === "seqFragment");

  // 1. Position participants horizontally (immutable — create new objects)
  const positionedParticipants = participants.map((p, i) => ({
    ...p,
    position: { x: i * SEQ_LAYOUT.participantSpacing, y: 0 },
  }));

  // 2. Compute lifeline height from message count
  const lifelineHeight =
    SEQ_LAYOUT.messageStartY +
    edges.length * SEQ_LAYOUT.messageSpacing +
    SEQ_LAYOUT.lifelinePadding;

  // 3. Enrich edges with order and messageY (immutable)
  const enrichedEdges = edges.map((edge, i) => ({
    ...edge,
    data: {
      ...edge.data,
      order: i,
      messageY: SEQ_LAYOUT.messageStartY + i * SEQ_LAYOUT.messageSpacing,
    },
  }));

  // 4. Compute activation bars per participant
  //    Activation: starts when a sync message arrives (target), ends when a return is sent (source)
  const activationsMap = new Map<
    string,
    { y: number; height: number }[]
  >();
  for (const p of positionedParticipants) {
    activationsMap.set(p.id, []);
  }

  // Track open activations per participant: stack of start Y values
  const openActivations = new Map<string, number[]>();
  for (const p of positionedParticipants) {
    openActivations.set(p.id, []);
  }

  for (const edge of enrichedEdges) {
    const edgeData = edge.data as Record<string, unknown>;
    const msgY = edgeData.messageY as number;
    const originalData = (edge.data as Record<string, unknown> | undefined);
    const edgeType = (originalData?.type as string) ?? "sync";

    if (edgeType === "sync" || edgeType === "async") {
      // Open activation on the target participant when sync/async message received
      const targetStack = openActivations.get(edge.target);
      if (targetStack) {
        targetStack.push(msgY);
      }
    } else if (edgeType === "return") {
      // Close activation on the source participant when return is sent
      const sourceStack = openActivations.get(edge.source);
      if (sourceStack && sourceStack.length > 0) {
        const startY = sourceStack.pop()!;
        const acts = activationsMap.get(edge.source) ?? [];
        acts.push({ y: startY, height: msgY - startY });
        activationsMap.set(edge.source, acts);
      }
    }
  }

  // Close any remaining open activations at the end of the lifeline
  for (const [participantId, stack] of openActivations) {
    const acts = activationsMap.get(participantId) ?? [];
    for (const startY of stack) {
      acts.push({
        y: startY,
        height: lifelineHeight - SEQ_LAYOUT.lifelinePadding - startY,
      });
    }
    activationsMap.set(participantId, acts);
  }

  // 5. Set participant height and activation data (immutable)
  const enrichedParticipants = positionedParticipants.map((p) => ({
    ...p,
    style: { ...p.style, height: lifelineHeight },
    data: {
      ...p.data,
      lifelineHeight,
      activations: activationsMap.get(p.id) ?? [],
    },
  }));

  // 6. Build participant X lookup map
  const participantXMap = new Map(
    enrichedParticipants.map((p) => [p.id, p.position.x]),
  );

  // 7. Position fragment nodes around their messages (immutable)
  const positionedFragments = fragments.map((frag) => {
    const fragData = frag.data as Record<string, unknown>;
    const messageIds = ((fragData.group as string) || "")
      .split(",")
      .map((s) => s.trim())
      .filter(Boolean);
    const messageIndices = messageIds
      .map((id) => enrichedEdges.findIndex((e) => e.id === id))
      .filter((i) => i >= 0);

    if (messageIndices.length === 0) {
      return { ...frag };
    }

    const minIdx = Math.min(...messageIndices);
    const maxIdx = Math.max(...messageIndices);
    const topY =
      SEQ_LAYOUT.messageStartY +
      minIdx * SEQ_LAYOUT.messageSpacing -
      SEQ_LAYOUT.fragmentPadding;
    const bottomY =
      SEQ_LAYOUT.messageStartY +
      maxIdx * SEQ_LAYOUT.messageSpacing +
      SEQ_LAYOUT.fragmentPadding;

    // Find leftmost and rightmost participants involved in fragment messages
    const involvedEdges = messageIndices.map((i) => enrichedEdges[i]!);
    const allParticipantIds = new Set(
      involvedEdges.flatMap((e) => [e.source, e.target]),
    );
    const xPositions = [...allParticipantIds].map(
      (id) => participantXMap.get(id) ?? 0,
    );
    const minX = Math.min(...xPositions) - SEQ_LAYOUT.fragmentPadding;
    const maxX =
      Math.max(...xPositions) +
      SEQ_SIZES.participant.w +
      SEQ_LAYOUT.fragmentPadding;

    return {
      ...frag,
      position: { x: minX, y: topY },
      style: {
        width: maxX - minX,
        height: bottomY - topY,
      },
    };
  });

  return {
    nodes: [...enrichedParticipants, ...positionedFragments],
    edges: enrichedEdges,
  };
}
