export type TraceFlowPair = { fromId: string; toId: string };

export type RoutedTraceFlowEdge = TraceFlowPair & {
  sourceHandle: string;
  targetHandle: string;
  offset: number;
  pathMode?: "smooth" | "corridor";
  corridorY?: number;
  /** Added to React Flow sourceX for the first horizontal stub. */
  corridorExitOffset?: number;
  /** Added to React Flow targetX for the vertical drop into the target. */
  corridorEntryOffset?: number;
  /**
   * Added to React Flow sourceX for the late branch point on the corridor.
   * Shared-trunk fans travel together until this X, then fork to each target.
   */
  corridorBranchOffset?: number;
};

export type TraceFlowNodeRect = {
  x: number;
  y: number;
  width: number;
  height: number;
  laneIndex: number;
};

export type TraceFlowLaneLayout = {
  laneTops: number[];
  laneHeights: number[];
  laneGap: number;
};

const OFFSET_STEP = 18;
const CORRIDOR_STEP = 14;
const HORIZONTAL_STUB = 14;
const MIN_HORIZONTAL_SPAN = 24;

export const TRACE_FLOW_HANDLE_IN = "in";
export const TRACE_FLOW_HANDLE_OUT = "out";

const stableSortPairs = (list: TraceFlowPair[]): TraceFlowPair[] =>
  [...list].sort((a, b) => {
    const keyA = `${a.fromId}->${a.toId}`;
    const keyB = `${b.fromId}->${b.toId}`;
    return keyA.localeCompare(keyB);
  });

const edgeKey = (fromId: string, toId: string) => `${fromId}->${toId}`;

const targetColumnKey = (rect: TraceFlowNodeRect) =>
  `${rect.laneIndex}:${Math.round(rect.x)}`;

const sourceColumnKey = (rect: TraceFlowNodeRect) =>
  `${rect.laneIndex}:${Math.round(rect.x)}`;

type CorridorMeta = {
  corridorY: number;
  corridorExitOffset: number;
  corridorEntryOffset: number;
  corridorBranchOffset?: number;
};

/** Spread multiple edges from the same node via handles + smoothstep offset. */
export const assignTraceFlowEdgeRouting = (
  pairs: TraceFlowPair[],
): {
  routed: RoutedTraceFlowEdge[];
  incomingCount: Map<string, number>;
  outgoingCount: Map<string, number>;
} => {
  const bySource = new Map<string, TraceFlowPair[]>();
  const byTarget = new Map<string, TraceFlowPair[]>();

  pairs.forEach((p) => {
    const outList = bySource.get(p.fromId) ?? [];
    outList.push(p);
    bySource.set(p.fromId, outList);

    const inList = byTarget.get(p.toId) ?? [];
    inList.push(p);
    byTarget.set(p.toId, inList);
  });

  bySource.forEach((list, id) => bySource.set(id, stableSortPairs(list)));
  byTarget.forEach((list, id) => byTarget.set(id, stableSortPairs(list)));

  const incomingCount = new Map<string, number>();
  const outgoingCount = new Map<string, number>();
  byTarget.forEach((list, id) => incomingCount.set(id, list.length));
  // 1→N fans use a single source handle (shared trunk); report 1 outbound slot.
  bySource.forEach((list, id) => outgoingCount.set(id, list.length > 1 ? 1 : list.length));

  const routed: RoutedTraceFlowEdge[] = pairs.map((p) => {
    const outList = bySource.get(p.fromId) ?? [p];
    const inList = byTarget.get(p.toId) ?? [p];
    const sourceIndex = outList.findIndex((e) => e.fromId === p.fromId && e.toId === p.toId);
    const targetIndex = inList.findIndex((e) => e.fromId === p.fromId && e.toId === p.toId);

    // 1→N fans leave from one handle so the stroke starts as a single trunk.
    const sourceHandle =
      outList.length > 1
        ? `${TRACE_FLOW_HANDLE_OUT}-0`
        : `${TRACE_FLOW_HANDLE_OUT}-${Math.max(0, sourceIndex)}`;
    const targetHandle = `${TRACE_FLOW_HANDLE_IN}-${Math.max(0, targetIndex)}`;

    let offset = 0;
    if (outList.length > 1) {
      offset = 0;
    } else if (inList.length > 1) {
      offset = (targetIndex - (inList.length - 1) / 2) * OFFSET_STEP;
    }

    return { ...p, sourceHandle, targetHandle, offset, pathMode: "smooth" };
  });

  return { routed, incomingCount, outgoingCount };
};

const baseCorridorY = (
  src: TraceFlowNodeRect,
  srcLane: number,
  tgtLane: number,
  lanes: TraceFlowLaneLayout,
): number => {
  if (srcLane !== tgtLane) {
    const srcLaneBottom =
      (lanes.laneTops[srcLane] ?? src.y) + (lanes.laneHeights[srcLane] ?? 0);
    return srcLaneBottom + lanes.laneGap * 0.42;
  }
  return src.y + src.height + 12;
};

/**
 * Route long horizontal spans through the lane gutter so edges do not cut across
 * unrelated nodes in the same swimlane.
 * 1→N source fans share one corridor trunk and only branch near targets.
 */
export const enrichTraceFlowCorridorRouting = (
  routed: RoutedTraceFlowEdge[],
  rects: Map<string, TraceFlowNodeRect>,
  lanes: TraceFlowLaneLayout,
): RoutedTraceFlowEdge[] => {
  const outDegree = new Map<string, number>();
  routed.forEach((edge) => {
    outDegree.set(edge.fromId, (outDegree.get(edge.fromId) ?? 0) + 1);
  });

  const corridorCandidates: RoutedTraceFlowEdge[] = [];

  routed.forEach((edge) => {
    const src = rects.get(edge.fromId);
    const tgt = rects.get(edge.toId);
    if (!src || !tgt) return;

    const srcRight = src.x + src.width;
    const crossLane = src.laneIndex !== tgt.laneIndex;
    const longSpan = tgt.x > srcRight + MIN_HORIZONTAL_SPAN;
    const sourceFan = (outDegree.get(edge.fromId) ?? 0) > 1;
    if (crossLane || longSpan || sourceFan) {
      corridorCandidates.push(edge);
    }
  });

  const corridorMetaByKey = new Map<string, CorridorMeta>();
  const groups = new Map<string, RoutedTraceFlowEdge[]>();

  corridorCandidates.forEach((edge) => {
    const src = rects.get(edge.fromId);
    const tgt = rects.get(edge.toId);
    if (!src || !tgt) return;
    const groupKey = `${src.laneIndex}->${tgt.laneIndex}`;
    const list = groups.get(groupKey) ?? [];
    list.push(edge);
    groups.set(groupKey, list);
  });

  groups.forEach((edges, groupKey) => {
    const [srcLaneStr, tgtLaneStr] = groupKey.split("->");
    const srcLane = Number(srcLaneStr);
    const tgtLane = Number(tgtLaneStr);

    const byFrom = new Map<string, RoutedTraceFlowEdge[]>();
    edges.forEach((edge) => {
      const list = byFrom.get(edge.fromId) ?? [];
      list.push(edge);
      byFrom.set(edge.fromId, list);
    });

    // Distinct sources in this lane-pair get separated corridor Y bands.
    const sourceIds = Array.from(byFrom.keys()).sort((a, b) => {
      const ax = rects.get(a)?.x ?? 0;
      const bx = rects.get(b)?.x ?? 0;
      if (ax !== bx) return ax - bx;
      const ay = rects.get(a)?.y ?? 0;
      const by = rects.get(b)?.y ?? 0;
      return ay - by || a.localeCompare(b);
    });

    sourceIds.forEach((fromId, sourceIndex) => {
      const fan = [...(byFrom.get(fromId) ?? [])].sort((a, b) => {
        const aty = rects.get(a.toId)?.y ?? 0;
        const bty = rects.get(b.toId)?.y ?? 0;
        if (aty !== bty) return aty - bty;
        const atx = rects.get(a.toId)?.x ?? 0;
        const btx = rects.get(b.toId)?.x ?? 0;
        return (
          atx - btx ||
          edgeKey(a.fromId, a.toId).localeCompare(edgeKey(b.fromId, b.toId))
        );
      });
      const src = rects.get(fromId);
      if (!src || !fan.length) return;

      const baseY = baseCorridorY(src, srcLane, tgtLane, lanes);
      const corridorY = baseY + sourceIndex * CORRIDOR_STEP;
      const sharedExit = HORIZONTAL_STUB;
      const srcRight = src.x + src.width;
      const isFan = fan.length > 1;
      const entrySpread = isFan ? (fan.length - 1) * CORRIDOR_STEP : 0;

      const fanMetas = fan.map((edge, fanIndex) => {
        const tgt = rects.get(edge.toId);
        const corridorEntryOffset = isFan
          ? -(HORIZONTAL_STUB + entrySpread) + fanIndex * CORRIDOR_STEP
          : -HORIZONTAL_STUB;
        const entryXLayout = (tgt?.x ?? 0) + corridorEntryOffset;
        return { edge, corridorEntryOffset, entryXLayout };
      });

      let sharedBranchOffset: number | undefined;
      if (isFan && fanMetas.length) {
        const minEntryX = Math.min(...fanMetas.map((m) => m.entryXLayout));
        sharedBranchOffset = minEntryX - srcRight;
      }

      fanMetas.forEach(({ edge, corridorEntryOffset }) => {
        corridorMetaByKey.set(edgeKey(edge.fromId, edge.toId), {
          corridorY,
          corridorExitOffset: sharedExit,
          corridorEntryOffset,
          corridorBranchOffset: sharedBranchOffset,
        });
      });
    });
  });

  // Separate vertical drops into the same target column (stacked nodes),
  // without breaking shared-trunk exit / branch for same-source fans.
  const byTargetColumn = new Map<string, RoutedTraceFlowEdge[]>();
  corridorCandidates.forEach((edge) => {
    const tgt = rects.get(edge.toId);
    if (!tgt) return;
    const key = targetColumnKey(tgt);
    const list = byTargetColumn.get(key) ?? [];
    list.push(edge);
    byTargetColumn.set(key, list);
  });

  byTargetColumn.forEach((edges) => {
    if (edges.length <= 1) return;

    const sorted = [...edges].sort((a, b) => {
      const aty = rects.get(a.toId)?.y ?? 0;
      const bty = rects.get(b.toId)?.y ?? 0;
      return aty - bty || edgeKey(a.fromId, a.toId).localeCompare(edgeKey(b.fromId, b.toId));
    });

    const spread = (sorted.length - 1) * CORRIDOR_STEP;

    sorted.forEach((edge, index) => {
      const key = edgeKey(edge.fromId, edge.toId);
      const meta = corridorMetaByKey.get(key);
      if (!meta) return;
      // Same-source fans already have entry offsets relative to the fan.
      if ((outDegree.get(edge.fromId) ?? 0) > 1) return;
      corridorMetaByKey.set(key, {
        ...meta,
        corridorEntryOffset: -(HORIZONTAL_STUB + spread) + index * CORRIDOR_STEP,
      });
    });
  });

  // Separate trunks when multiple distinct sources share a column (stacked nodes).
  // Same-source fan edges keep a shared exit; different sources get staggered exits.
  const bySourceColumn = new Map<string, string[]>();
  corridorCandidates.forEach((edge) => {
    const src = rects.get(edge.fromId);
    if (!src) return;
    const key = sourceColumnKey(src);
    const list = bySourceColumn.get(key) ?? [];
    if (!list.includes(edge.fromId)) list.push(edge.fromId);
    bySourceColumn.set(key, list);
  });

  bySourceColumn.forEach((sourceIds) => {
    if (sourceIds.length <= 1) return;

    const sortedSources = [...sourceIds].sort((a, b) => {
      const asy = rects.get(a)?.y ?? 0;
      const bsy = rects.get(b)?.y ?? 0;
      return asy - bsy || a.localeCompare(b);
    });

    const spread = (sortedSources.length - 1) * CORRIDOR_STEP;

    sortedSources.forEach((fromId, index) => {
      const exitOffset = HORIZONTAL_STUB - spread / 2 + index * CORRIDOR_STEP;
      corridorCandidates
        .filter((e) => e.fromId === fromId)
        .forEach((edge) => {
          const key = edgeKey(edge.fromId, edge.toId);
          const meta = corridorMetaByKey.get(key);
          if (!meta) return;
          corridorMetaByKey.set(key, {
            ...meta,
            corridorExitOffset: exitOffset,
          });
        });
    });
  });

  return routed.map((edge) => {
    const meta = corridorMetaByKey.get(edgeKey(edge.fromId, edge.toId));
    if (!meta) return edge;
    return {
      ...edge,
      pathMode: "corridor",
      corridorY: meta.corridorY,
      corridorExitOffset: meta.corridorExitOffset,
      corridorEntryOffset: meta.corridorEntryOffset,
      corridorBranchOffset: meta.corridorBranchOffset,
      // Shared trunk fans already leave from out-0 with zero smooth offset.
      offset: (outDegree.get(edge.fromId) ?? 0) > 1 ? 0 : edge.offset,
      sourceHandle:
        (outDegree.get(edge.fromId) ?? 0) > 1
          ? `${TRACE_FLOW_HANDLE_OUT}-0`
          : edge.sourceHandle,
    };
  });
};

/**
 * Orthogonal path: exit source → drop to corridor → travel →
 * (optional late branch) → rise to target.
 */
export const buildTraceFlowCorridorPath = (
  sourceX: number,
  sourceY: number,
  targetX: number,
  targetY: number,
  corridorY: number,
  corridorEntryOffset?: number,
  corridorExitOffset?: number,
  corridorBranchOffset?: number,
): string => {
  const exitX = sourceX + (corridorExitOffset ?? HORIZONTAL_STUB);
  const entryX = targetX + (corridorEntryOffset ?? -HORIZONTAL_STUB);
  const midY = corridorY;
  const branchX =
    corridorBranchOffset != null ? sourceX + corridorBranchOffset : entryX;

  // Keep branch between exit and entry so the trunk never backtracks oddly.
  const clampedBranchX = Math.min(
    Math.max(branchX, Math.min(exitX, entryX)),
    Math.max(exitX, entryX),
  );

  if (Math.abs(clampedBranchX - entryX) < 0.5) {
    return [
      `M ${sourceX} ${sourceY}`,
      `L ${exitX} ${sourceY}`,
      `L ${exitX} ${midY}`,
      `L ${entryX} ${midY}`,
      `L ${entryX} ${targetY}`,
      `L ${targetX} ${targetY}`,
    ].join(" ");
  }

  return [
    `M ${sourceX} ${sourceY}`,
    `L ${exitX} ${sourceY}`,
    `L ${exitX} ${midY}`,
    `L ${clampedBranchX} ${midY}`,
    `L ${entryX} ${midY}`,
    `L ${entryX} ${targetY}`,
    `L ${targetX} ${targetY}`,
  ].join(" ");
};

/** Vertical position (%) for handle index among `count` siblings on one node side. */
export const traceFlowHandleTopPercent = (index: number, count: number): string => {
  if (count <= 1) return "50%";
  return `${((index + 1) / (count + 1)) * 100}%`;
};
