ChatApp/apps/desktop/src/components/MessageList.tsx

import { useVirtualizer } from '@tanstack/react-virtual';
import {
  forwardRef,
  useCallback,
  useEffect,
  useImperativeHandle,
  useLayoutEffect,
  useRef,
  useState,
  type ReactNode,
} from 'react';

import { isNearBottom, isNearTop, nextStickIntent } from '../lib/scrollController';
import type { VirtuosoRow } from '../pages/ConversationPage';

export interface MessageListHandle {
  scrollToBottom(behavior?: ScrollBehavior): void;
  scrollToRow(index: number, align?: 'center' | 'end', behavior?: ScrollBehavior): void;
}

export interface MessageListProps {
  rows: VirtuosoRow[];
  renderRow: (index: number, row: VirtuosoRow) => ReactNode;
  computeKey: (row: VirtuosoRow) => string;
  /** Initial scroll target for a freshly-mounted list. */
  initialAnchor: { type: 'bottom' } | { type: 'row'; index: number };
  /** Reveal gate — the list stays hidden until reactions/heights are loaded, so
   *  the post-paint height cascade is never visible. */
  ready: boolean;
  estimateRowHeight?: number;
  atBottomThreshold?: number;
  onReachTop?: () => void;
  onAtBottomChange?: (atBottom: boolean) => void;
  onTopRowChange?: (topIndex: number) => void;
}

export const MessageList = forwardRef<MessageListHandle, MessageListProps>(function MessageList(
  {
    rows,
    renderRow,
    computeKey,
    initialAnchor,
    ready,
    estimateRowHeight = 64,
    atBottomThreshold = 64,
    onReachTop,
    onAtBottomChange,
    onTopRowChange,
  },
  ref,
) {
  const scrollElRef = useRef<HTMLDivElement>(null);
  const [revealed, setRevealed] = useState(false);

  // THE single source of truth: should the view stay pinned to the bottom?
  // Only a genuine user up-input (wheel / key / touch / scrollbar drag) turns
  // this OFF; only reaching the bottom turns it ON. A measurement reflow must
  // never flip it — that was the root cause of the chat-switch bug.
  const stickRef = useRef(true);
  // Debounce for onAtBottomChange — fire the parent only on a real transition.
  const lastReportedAtBottomRef = useRef<boolean | null>(null);
  // Guard: scrolls WE cause (pin / measure re-pin / scrollToIndex) fire onScroll
  // a tick later. Within this window we don't treat a scrollTop decrease as the
  // user dragging up.
  const programmaticRef = useRef(0);
  // Previous scrollTop, to detect a genuine scrollbar/keyboard up-drag.
  const lastScrollTopRef = useRef(0);
  // Load-older preservation: remember the first row key + scrollHeight so a
  // prepend can be detected and the viewport restored.
  const prevFirstKeyRef = useRef<string | null>(null);
  const prevScrollHeightRef = useRef(0);

  // Latest onAtBottomChange, read through a ref so the input-listener effect
  // can stay mounted once (deps []) without capturing a stale callback.
  const onAtBottomChangeRef = useRef(onAtBottomChange);
  onAtBottomChangeRef.current = onAtBottomChange;

  const virtualizer = useVirtualizer({
    count: rows.length,
    getScrollElement: () => scrollElRef.current,
    estimateSize: () => estimateRowHeight,
    overscan: 8,
    getItemKey: (index) => computeKey(rows[index]!),
  });

  const readMetrics = useCallback(() => {
    const el = scrollElRef.current;
    return el
      ? { scrollTop: el.scrollTop, scrollHeight: el.scrollHeight, clientHeight: el.clientHeight }
      : { scrollTop: 0, scrollHeight: 0, clientHeight: 0 };
  }, []);

  const pinToBottom = useCallback(() => {
    const el = scrollElRef.current;
    if (!el) return;
    programmaticRef.current = performance.now();
    el.scrollTop = el.scrollHeight;
    lastScrollTopRef.current = el.scrollTop;
  }, []);

  // Report at-bottom to the parent only on a true transition, always driven by
  // the INTENT (stickRef) — never the raw position. This is what kills the
  // feedback loop: a transient "not at bottom" mid-reflow is never persisted.
  const reportAtBottom = useCallback((atBottom: boolean) => {
    if (lastReportedAtBottomRef.current === atBottom) return;
    lastReportedAtBottomRef.current = atBottom;
    onAtBottomChangeRef.current?.(atBottom);
  }, []);

  // A genuine user up-input: drop the stick intent immediately.
  const markUserMovedUp = useCallback(() => {
    if (!stickRef.current) return;
    stickRef.current = false;
    reportAtBottom(false);
  }, [reportAtBottom]);

  // Re-pin to the true bottom whenever the content (or viewport) resizes while
  // sticking. ResizeObserver fires after layout / before paint, so as rows
  // measure and the list grows the bottom stays pinned with no stale frame.
  useEffect(() => {
    const el = scrollElRef.current;
    if (!el) return;
    const ro = new ResizeObserver(() => {
      const e = scrollElRef.current;
      if (stickRef.current && e) {
        programmaticRef.current = performance.now();
        e.scrollTop = e.scrollHeight;
        lastScrollTopRef.current = e.scrollTop;
      }
    });
    ro.observe(el);
    const inner = el.firstElementChild;
    if (inner) ro.observe(inner);
    return () => ro.disconnect();
  }, []);

  // Genuine-user-intent listeners. These are the ONLY way (besides reaching the
  // bottom) the stick intent turns off, so a reflow can never unstick the list.
  useEffect(() => {
    const el = scrollElRef.current;
    if (!el) return;
    const onWheel = (e: WheelEvent) => {
      if (e.deltaY < 0) markUserMovedUp();
    };
    const onKeyDown = (e: KeyboardEvent) => {
      if (e.key === 'PageUp' || e.key === 'Home' || e.key === 'ArrowUp') markUserMovedUp();
    };
    let touchStartY = 0;
    const onTouchStart = (e: TouchEvent) => {
      touchStartY = e.touches[0]?.clientY ?? 0;
    };
    const onTouchMove = (e: TouchEvent) => {
      const y = e.touches[0]?.clientY ?? 0;
      // Finger dragged DOWN (content scrolls up toward older messages). Guard on
      // scrollTop>0 so an overscroll bounce at the bottom doesn't unstick.
      if (y - touchStartY > 8 && (scrollElRef.current?.scrollTop ?? 0) > 0) markUserMovedUp();
    };
    el.addEventListener('wheel', onWheel, { passive: true });
    el.addEventListener('keydown', onKeyDown);
    el.addEventListener('touchstart', onTouchStart, { passive: true });
    el.addEventListener('touchmove', onTouchMove, { passive: true });
    return () => {
      el.removeEventListener('wheel', onWheel);
      el.removeEventListener('keydown', onKeyDown);
      el.removeEventListener('touchstart', onTouchStart);
      el.removeEventListener('touchmove', onTouchMove);
    };
  }, [markUserMovedUp]);

  // Deferred reveal: when ready, pin to the anchor and keep pinning each frame
  // until the list height has SETTLED over two consecutive frames, THEN reveal —
  // so what appears is already at its final position with no top-then-jump.
  useLayoutEffect(() => {
    if (!ready || revealed || rows.length === 0) return;
    const el = scrollElRef.current;
    if (!el) return;
    const rowIdx = Math.max(
      0,
      Math.min(initialAnchor.type === 'row' ? initialAnchor.index : 0, rows.length - 1),
    );
    if (initialAnchor.type === 'bottom') {
      stickRef.current = true;
      pinToBottom();
    } else {
      stickRef.current = false;
      programmaticRef.current = performance.now();
      virtualizer.scrollToIndex(rowIdx, { align: 'start' });
    }
    reportAtBottom(stickRef.current);

    let prevSH = -1;
    let stableFrames = 0;
    const settle = (attempts: number): void => {
      const e = scrollElRef.current;
      if (!e) {
        setRevealed(true);
        return;
      }
      programmaticRef.current = performance.now();
      if (stickRef.current) {
        e.scrollTop = e.scrollHeight;
        lastScrollTopRef.current = e.scrollTop;
      } else {
        virtualizer.scrollToIndex(rowIdx, { align: 'start' });
      }
      const sh = e.scrollHeight;
      // Require TWO consecutive stable-height frames: a single stable frame can
      // land mid-cascade (between reactions and the unread divider measuring)
      // and reveal a not-yet-final layout that then jumps.
      stableFrames = sh === prevSH ? stableFrames + 1 : 0;
      prevSH = sh;
      if (stableFrames >= 2 || attempts <= 0) {
        setRevealed(true);
      } else {
        requestAnimationFrame(() => settle(attempts - 1));
      }
    };
    requestAnimationFrame(() => settle(12));
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [ready, rows.length]);

  // Load-older preservation: if rows were prepended (first key changed and the
  // user is near the top), restore scrollTop by the height delta so the viewport
  // stays put instead of jumping.
  useLayoutEffect(() => {
    const firstKey = rows.length > 0 ? computeKey(rows[0]!) : null;
    const el = scrollElRef.current;
    if (el && revealed && prevFirstKeyRef.current && firstKey !== prevFirstKeyRef.current) {
      const delta = el.scrollHeight - prevScrollHeightRef.current;
      if (delta > 0 && el.scrollTop < atBottomThreshold * 4) {
        el.scrollTop += delta;
        lastScrollTopRef.current = el.scrollTop;
      }
    }
    prevFirstKeyRef.current = firstKey;
    prevScrollHeightRef.current = el?.scrollHeight ?? 0;
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [rows]);

  // Re-pin on any rows change while sticking. Covers the two-phase data swap
  // (the cached array is replaced by the freshly-decrypted one ~100ms after
  // reveal) which the ResizeObserver can miss when the new content happens to
  // measure to the same height.
  useLayoutEffect(() => {
    if (revealed && stickRef.current) pinToBottom();
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [rows]);

  const handleScroll = useCallback(() => {
    const m = readMetrics();
    const programmatic = performance.now() - programmaticRef.current < 120;
    const nearBottom = isNearBottom(m, atBottomThreshold);
    // A scrollbar drag or keyboard scroll surfaces here as a scrollTop decrease.
    // Suppress it inside the programmatic window so our own re-pin / settle is
    // never mistaken for the user moving up. 2px deadzone absorbs sub-pixel jitter.
    const userMovedUp = !programmatic && m.scrollTop < lastScrollTopRef.current - 2;
    lastScrollTopRef.current = m.scrollTop;

    stickRef.current = nextStickIntent(stickRef.current, { nearBottom, userMovedUp });
    reportAtBottom(stickRef.current);

    if (isNearTop(m, atBottomThreshold * 4)) onReachTop?.();
    const first = virtualizer.getVirtualItems()[0];
    if (first) onTopRowChange?.(first.index);
  }, [atBottomThreshold, onReachTop, onTopRowChange, readMetrics, reportAtBottom, virtualizer]);

  useImperativeHandle(
    ref,
    () => ({
      scrollToBottom: () => {
        stickRef.current = true;
        reportAtBottom(true);
        pinToBottom();
      },
      scrollToRow: (index, align = 'center') => {
        // The user is jumping to a specific row — drop the stick intent first so
        // the ResizeObserver doesn't immediately drag the target back to the bottom.
        stickRef.current = false;
        reportAtBottom(false);
        programmaticRef.current = performance.now();
        virtualizer.scrollToIndex(index, { align });
      },
    }),
    // eslint-disable-next-line react-hooks/exhaustive-deps
    [virtualizer, rows.length, reportAtBottom, pinToBottom],
  );

  const items = virtualizer.getVirtualItems();

  return (
    <div
      ref={scrollElRef}
      onScroll={handleScroll}
      tabIndex={0}
      className="min-h-0 flex-1 overflow-y-auto"
      style={{
        opacity: revealed ? 1 : 0,
        position: 'relative',
        overflowAnchor: 'none',
        outline: 'none',
      }}
    >
      <div style={{ height: virtualizer.getTotalSize(), position: 'relative', width: '100%' }}>
        {items.map((vi) => (
          <div
            key={vi.key}
            data-index={vi.index}
            ref={virtualizer.measureElement}
            style={{
              position: 'absolute',
              top: 0,
              left: 0,
              width: '100%',
              transform: `translateY(${vi.start}px)`,
            }}
          >
            {renderRow(vi.index, rows[vi.index]!)}
          </div>
        ))}
      </div>
      {/* 12px bottom breathing space (matches the old Virtuoso Footer). */}
      <div style={{ height: 12 }} />
    </div>
  );
});