ujsx/src/host/reconcile.ts

import { effect } from "@preact/signals-core";
import { Value } from "@alkdev/typebox/value";
import type { Fiber, Effect } from "./fiber.js";
import type { HostConfig } from "./config.js";
import type { UNode, UElement } from "../core/schema.js";
import { isUPrimitive, isURoot, isUElement } from "../core/schema.js";

export interface CommitContext<I> {
  host: HostConfig<string, I, unknown>;
  ctx: unknown;
  createFiber: (node: UNode, parentFiber: Fiber<I>) => Fiber<I>;
}

interface PendingUpdate {
  fiber: Fiber<unknown>;
  nextNode: UNode;
}

let pendingUpdates: PendingUpdate[] = [];
let flushScheduled = false;

export function scheduleUpdate<I>(
  fiber: Fiber<I>,
  nextNode: UNode,
  host: HostConfig<string, I, unknown>,
  ctx: unknown,
): void {
  pendingUpdates.push({ fiber: fiber as Fiber<unknown>, nextNode });
  if (!flushScheduled) {
    flushScheduled = true;
    queueMicrotask(() => flushUpdates(host as HostConfig<string, unknown, unknown>, ctx));
  }
}

export function flushUpdates(
  host: HostConfig<string, unknown, unknown>,
  ctx: unknown,
): void {
  flushScheduled = false;
  const updates = pendingUpdates.splice(0);

  for (const { fiber, nextNode } of updates) {
    reconcileProps(fiber, nextNode, host, ctx);
  }

  const seen = new Set<Fiber<unknown>>();
  for (const { fiber } of updates) {
    let root: Fiber<unknown> | null = fiber;
    while (root.parent) root = root.parent;
    if (!seen.has(root)) {
      seen.add(root);
      commitEffects(root, host, ctx);
    }
  }
}

export function reconcileProps<I>(
  fiber: Fiber<I>,
  nextNode: UNode,
  host: HostConfig<string, I, unknown>,
  ctx: unknown,
): void {
  if (fiber.cachedNode !== null && Value.Equal(fiber.cachedNode, nextNode)) {
    return;
  }

  if (isUPrimitive(nextNode)) {
    if (fiber.tag === "#text") {
      const text = nextNode === null ? "" : String(nextNode);
      if (fiber.props.text !== text) {
        const nextProps = { text };
        const payload = host.prepareUpdate?.(
          fiber.instance,
          fiber.tag as never,
          fiber.props,
          nextProps,
          ctx as never,
        );
        if (payload !== null && payload !== undefined) {
          fiber.effect = { type: "update", payload };
          fiber.prevProps = { ...fiber.props };
          fiber.props = nextProps;
        }
      }
    }
    fiber.cachedNode = nextNode;
    return;
  }

  if (isURoot(nextNode)) {
    const rootChildren = nextNode.children;
    const count = Math.min(fiber.children.length, rootChildren.length);
    for (let i = 0; i < count; i++) {
      reconcileProps(fiber.children[i]!, rootChildren[i]!, host, ctx);
    }
    fiber.cachedNode = nextNode;
    return;
  }

  const el = nextNode as UElement;

  if (typeof el.type === "function") {
    const component = el.type as (props: Record<string, unknown>) => UNode;
    const out = component({ ...el.props, children: el.children });
    reconcileProps(fiber, out, host, ctx);
    fiber.cachedNode = nextNode;
    return;
  }

  if (fiber.tag !== "#text" && fiber.tag !== el.type) return;

  const nextProps = el.props as Record<string, unknown>;
  const payload = host.prepareUpdate?.(
    fiber.instance,
    fiber.tag as never,
    fiber.props,
    nextProps,
    ctx as never,
  );

  if (payload !== null && payload !== undefined) {
    fiber.effect = { type: "update", payload };
    fiber.prevProps = { ...fiber.props };
    fiber.props = nextProps;
  }

  const count = Math.min(fiber.children.length, el.children.length);
  for (let i = 0; i < count; i++) {
    reconcileProps(fiber.children[i]!, el.children[i]!, host, ctx);
  }
  fiber.cachedNode = nextNode;
}

export function commitEffects<I>(
  fiber: Fiber<I>,
  host: HostConfig<string, I, unknown>,
  ctx: unknown,
): void {
  if (fiber.effect?.type === "update") {
    const updateEffect = fiber.effect as Effect<I> & { type: "update"; payload: unknown };
    host.commitUpdate?.(
      fiber.instance,
      updateEffect.payload,
      fiber.tag as never,
      fiber.prevProps ?? {},
      fiber.props,
      ctx as never,
    );
  }
  for (const child of fiber.children) {
    commitEffects(child, host, ctx);
  }
  fiber.effect = null;
}

export function wireSignalToFiber<I>(
  fiber: Fiber<I>,
  signalGetter: () => UNode,
  host: HostConfig<string, I, unknown>,
  ctx: unknown,
): void {
  const disposer = effect(() => {
    const nextNode = signalGetter();
    scheduleUpdate(fiber, nextNode, host, ctx);
  });
  fiber.signalDisposers.push(disposer);
}

export function resetUpdateQueue(): void {
  pendingUpdates = [];
  flushScheduled = false;
}

export function longestIncreasingSubsequence(arr: number[]): number[] {
  if (arr.length === 0) return [];

  const tails: number[] = [];
  const tailsIdx: number[] = [];
  const prev: number[] = new Array(arr.length).fill(-1);

  for (let i = 0; i < arr.length; i++) {
    const val = arr[i]!;
    let lo = 0;
    let hi = tails.length;
    while (lo < hi) {
      const mid = (lo + hi) >>> 1;
      if (tails[mid]! < val) {
        lo = mid + 1;
      } else {
        hi = mid;
      }
    }
    if (lo > 0) {
      prev[i] = tailsIdx[lo - 1]!;
    }
    if (lo === tails.length) {
      tails.push(val);
      tailsIdx.push(i);
    } else {
      tails[lo] = val;
      tailsIdx[lo] = i;
    }
  }

  const result: number[] = new Array(tails.length);
  let k = tailsIdx[tails.length - 1]!;
  for (let j = tails.length - 1; j >= 0; j--) {
    result[j] = k;
    k = prev[k]!;
  }
  return result;
}

export interface MatchedChild<I> {
  oldFiber: Fiber<I>;
  newChild: UElement;
  index: number;
}

export interface ChildClassification<I> {
  matched: MatchedChild<I>[];
  added: { newChild: UElement; index: number }[];
  removed: Fiber<I>[];
  moves: Map<number, Fiber<I>>;
}

export function reconcileChildren<I>(
  oldFibers: Fiber<I>[],
  newChildren: UNode[],
): ChildClassification<I> {
  const matched: MatchedChild<I>[] = [];
  const added: { newChild: UElement; index: number }[] = [];
  const removed: Fiber<I>[] = [];

  const oldKeyMap = new Map<string, Fiber<I>>();
  const displacedOldFibers: Fiber<I>[] = [];
  for (const fiber of oldFibers) {
    if (fiber.key !== undefined) {
      if (oldKeyMap.has(fiber.key)) {
        console.warn(`Duplicate key "${fiber.key}" among old children; last-wins`);
        displacedOldFibers.push(oldKeyMap.get(fiber.key)!);
      }
      oldKeyMap.set(fiber.key, fiber);
    }
  }

  const matchedOldKeys = new Set<string>();
  const unkeyedOldUsed = new Set<number>();
  let unkeyedOldCursor = 0;

  const seenNewKeys = new Map<string, number>();

  for (let i = 0; i < newChildren.length; i++) {
    const child = newChildren[i]!;
    if (!isUElement(child)) continue;

    if (child.key !== undefined) {
      const prevIdx = seenNewKeys.get(child.key);
      if (prevIdx !== undefined) {
        console.warn(`Duplicate key "${child.key}" among new children; last-wins`);
        const prevMatch = matched.findIndex((m) => m.newChild.key === child.key && m.index === prevIdx);
        if (prevMatch !== -1) {
          const oldFiber = matched[prevMatch]!.oldFiber;
          removed.push(oldFiber);
          matched.splice(prevMatch, 1);
        }
      }
      seenNewKeys.set(child.key, i);

      const oldFiber = oldKeyMap.get(child.key);
      if (oldFiber !== undefined) {
        matchedOldKeys.add(child.key);
        if (oldFiber.tag === child.type) {
          matched.push({ oldFiber, newChild: child, index: i });
        } else {
          removed.push(oldFiber);
          added.push({ newChild: child, index: i });
        }
      } else {
        added.push({ newChild: child, index: i });
      }
    } else {
      let matchedOld: Fiber<I> | null = null;
      let matchedOldIdx = -1;
      for (let j = unkeyedOldCursor; j < oldFibers.length; j++) {
        if (!unkeyedOldUsed.has(j) && oldFibers[j]!.key === undefined) {
          matchedOld = oldFibers[j]!;
          matchedOldIdx = j;
          break;
        }
      }
      if (matchedOld !== null) {
        unkeyedOldUsed.add(matchedOldIdx);
        unkeyedOldCursor = matchedOldIdx + 1;
        if (matchedOld.tag === child.type) {
          matched.push({ oldFiber: matchedOld, newChild: child, index: i });
        } else {
          removed.push(matchedOld);
          added.push({ newChild: child, index: i });
        }
      } else {
        added.push({ newChild: child, index: i });
      }
    }
  }

  for (let i = 0; i < oldFibers.length; i++) {
    const fiber = oldFibers[i]!;
    if (fiber.key !== undefined) {
      if (!matchedOldKeys.has(fiber.key)) {
        removed.push(fiber);
      }
    } else {
      if (!unkeyedOldUsed.has(i)) {
        removed.push(fiber);
      }
    }
  }

  for (const displaced of displacedOldFibers) {
    if (matchedOldKeys.has(displaced.key!)) {
      removed.push(displaced);
    }
  }

  const oldIndices = matched.map((m) => oldFibers.indexOf(m.oldFiber));
  const lisPositions = longestIncreasingSubsequence(oldIndices);
  const lisSet = new Set(lisPositions);

  const moves = new Map<number, Fiber<I>>();
  for (let i = 0; i < matched.length; i++) {
    if (!lisSet.has(i)) {
      moves.set(i, matched[i]!.oldFiber);
    }
  }

  return { matched, added, removed, moves };
}

export function commitMutations<I>(
  parentFiber: Fiber<I>,
  classification: ChildClassification<I>,
  commitCtx: CommitContext<I>,
): void {
  const { host, ctx, createFiber } = commitCtx;
  const parentInst = parentFiber.instance;

  const movedSet = new Set<Fiber<I>>(classification.moves.values());

  // Map matched index (position in classification.matched[]) for quick lookup
  const matchedArrayIndex = new Map<MatchedChild<I>, number>();
  for (let i = 0; i < classification.matched.length; i++) {
    matchedArrayIndex.set(classification.matched[i]!, i);
  }

  // Build maps keyed by new-children position
  const matchedByNewIndex = new Map<number, MatchedChild<I>>();
  for (const m of classification.matched) {
    matchedByNewIndex.set(m.index, m);
  }
  const addedByIndex = new Map<number, UElement>();
  for (const { newChild, index } of classification.added) {
    addedByIndex.set(index, newChild);
  }

  // Collect all positions in new-children order
  const allIndices = new Set<number>();
  for (const m of classification.matched) allIndices.add(m.index);
  for (const a of classification.added) allIndices.add(a.index);
  const sortedIndices = [...allIndices].sort((a, b) => a - b);

  // Phase 1: Removes — reverse order (children before parents, bottom-up)
  for (let i = classification.removed.length - 1; i >= 0; i--) {
    const fiber = classification.removed[i]!;
    host.removeChild?.(parentInst as never, fiber.instance as never, ctx as never);
  }

  // Build new children array and determine placement actions
  const newChildren: Fiber<I>[] = [];
  type Placement = { fiber: Fiber<I>; beforeFiber: Fiber<I> | null; kind: "insert" | "move" };
  const placements: Placement[] = [];

  // Create fibers for added children first (so they're available for before-lookup)
  const addedFibers = new Map<number, Fiber<I>>();
  for (const { index } of classification.added) {
    const newChild = addedByIndex.get(index)!;
    const fiber = createFiber(newChild, parentFiber);
    addedFibers.set(index, fiber);
  }

  // Build the final children array
  for (const idx of sortedIndices) {
    const addedChild = addedByIndex.get(idx);
    if (addedChild !== undefined) {
      const fiber = addedFibers.get(idx)!;
      newChildren.push(fiber);
      // Find before: next fiber in newChildren that is "staying" (not moved, not just inserted)
      const before = findNextStayingFiber(sortedIndices, matchedByNewIndex, addedByIndex, movedSet, idx);
      placements.push({ fiber, beforeFiber: before, kind: "insert" });
      continue;
    }

    const match = matchedByNewIndex.get(idx);
    if (match !== undefined) {
      const fiber = match.oldFiber;
      newChildren.push(fiber);
      const mIdx = matchedArrayIndex.get(match)!;
      if (classification.moves.has(mIdx)) {
        const before = findNextStayingFiber(sortedIndices, matchedByNewIndex, addedByIndex, movedSet, idx);
        placements.push({ fiber, beforeFiber: before, kind: "move" });
      }
    }
  }

  // Phase 2: Inserts + Moves — left-to-right
  for (const { fiber, beforeFiber } of placements) {
    if (host.insertBefore && beforeFiber) {
      host.insertBefore(
        parentInst as never,
        fiber.instance as never,
        beforeFiber.instance as never,
        ctx as never,
      );
    } else {
      host.appendChild(parentInst as never, fiber.instance as never, ctx as never);
    }
  }

  // Update fiber tree
  parentFiber.children = newChildren;
  for (const fiber of classification.removed) {
    fiber.parent = null;
  }

  // Phase 3: Updates — top-down (parent before child) via commitEffects
  commitEffects(parentFiber, host as HostConfig<string, I, unknown>, ctx);
}

function findNextStayingFiber<I>(
  sortedIndices: number[],
  matchedByNewIndex: Map<number, MatchedChild<I>>,
  addedByIndex: Map<number, UElement>,
  movedSet: Set<Fiber<I>>,
  currentIdx: number,
): Fiber<I> | null {
  const posInSorted = sortedIndices.indexOf(currentIdx);
  for (let p = posInSorted + 1; p < sortedIndices.length; p++) {
    const nextIdx = sortedIndices[p]!;
    const isAdded = addedByIndex.has(nextIdx);
    const match = matchedByNewIndex.get(nextIdx);
    if (!isAdded && match !== undefined) {
      const isMoved = movedSet.has(match.oldFiber);
      if (!isMoved) {
        return match.oldFiber;
      }
    }
  }
  return null;
}