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

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

import { isNearBottom, isNearTop } 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 behind a spinner until true, so the
   *  post-paint height cascade (reactions/divider) 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);
  const atBottomRef = useRef(true);
  // Intent: keep the view pinned to the bottom? Only a genuine user scroll flips
  // this — measurement reflows must not.
  const stickRef = useRef(true);
  // Guard: scroll events we cause (pin / measure re-pin) fire a tick after we set
  // scrollTop. handleScroll ignores events within this window so a measurement
  // reflow is never mistaken for the user scrolling up.
  const programmaticRef = 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);

  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;
  }, []);

  // 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 — THIS is what
  // carries the list to the bottom through the measure settle (the old
  // getTotalSize effect lagged a frame, letting onScroll wrongly flip atBottom).
  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;
      }
    });
    ro.observe(el);
    const inner = el.firstElementChild;
    if (inner) ro.observe(inner);
    return () => ro.disconnect();
  }, []);

  // Deferred reveal: when ready, pin to the anchor and keep pinning each frame
  // until the list height has SETTLED (rows measure over 1-2 frames and grow it),
  // THEN reveal — so what appears is already at the final bottom, with no
  // top-then-jump flicker. The ResizeObserver above keeps it pinned afterwards.
  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;
      atBottomRef.current = true;
      pinToBottom();
    } else {
      stickRef.current = false;
      atBottomRef.current = false;
      virtualizer.scrollToIndex(rowIdx, { align: 'start' });
    }
    onAtBottomChange?.(atBottomRef.current);
    let prevSH = -1;
    const settle = (attempts: number): void => {
      const e = scrollElRef.current;
      if (!e) {
        setRevealed(true);
        return;
      }
      if (stickRef.current) {
        programmaticRef.current = performance.now();
        e.scrollTop = e.scrollHeight;
      } else {
        virtualizer.scrollToIndex(rowIdx, { align: 'start' });
      }
      const sh = e.scrollHeight;
      if (sh === prevSH || attempts <= 0) {
        setRevealed(true);
      } else {
        prevSH = sh;
        requestAnimationFrame(() => settle(attempts - 1));
      }
    };
    requestAnimationFrame(() => settle(10));
    // 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;
      }
    }
    prevFirstKeyRef.current = firstKey;
    prevScrollHeightRef.current = el?.scrollHeight ?? 0;
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [rows]);

  const handleScroll = useCallback(() => {
    // Ignore scroll events we triggered (pin / measure re-pin); they fire a tick
    // after we set scrollTop. Only a genuine user scroll updates the stick intent —
    // otherwise a measurement reflow wrongly flips atBottom and stops the pinning.
    if (performance.now() - programmaticRef.current < 120) return;
    const m = readMetrics();
    const atBottom = isNearBottom(m, atBottomThreshold);
    stickRef.current = atBottom;
    if (atBottom !== atBottomRef.current) {
      atBottomRef.current = atBottom;
      onAtBottomChange?.(atBottom);
    }
    if (isNearTop(m, atBottomThreshold * 4)) onReachTop?.();
    const first = virtualizer.getVirtualItems()[0];
    if (first) onTopRowChange?.(first.index);
  }, [atBottomThreshold, onAtBottomChange, onReachTop, onTopRowChange, readMetrics, virtualizer]);

  useImperativeHandle(
    ref,
    () => ({
      scrollToBottom: () => {
        stickRef.current = true;
        atBottomRef.current = true;
        pinToBottom();
      },
      scrollToRow: (index, align = 'center') => {
        virtualizer.scrollToIndex(index, { align });
      },
    }),
    // eslint-disable-next-line react-hooks/exhaustive-deps
    [virtualizer, rows.length],
  );

  const items = virtualizer.getVirtualItems();

  return (
    <>
      <ScrollDebugOverlay
        scrollElRef={scrollElRef}
        getTotal={() => virtualizer.getTotalSize()}
        atBottomRef={atBottomRef}
        revealed={revealed}
        ready={ready}
      />
      <div
        ref={scrollElRef}
        onScroll={handleScroll}
        className="min-h-0 flex-1 overflow-y-auto"
        style={{ opacity: revealed ? 1 : 0, position: 'relative', overflowAnchor: '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>
    </>
  );
});

interface ScrollDebugOverlayProps {
  scrollElRef: RefObject<HTMLDivElement>;
  getTotal: () => number;
  atBottomRef: MutableRefObject<boolean>;
  revealed: boolean;
  ready: boolean;
}

// TEMP on-screen debug readout — remove after the chat-switch scroll bug is fixed.
function ScrollDebugOverlay({
  scrollElRef,
  getTotal,
  atBottomRef,
  revealed,
  ready,
}: ScrollDebugOverlayProps) {
  const [, force] = useState(0);
  useEffect(() => {
    let raf = 0;
    const loop = () => {
      force((n) => (n + 1) % 1_000_000);
      raf = requestAnimationFrame(loop);
    };
    raf = requestAnimationFrame(loop);
    return () => cancelAnimationFrame(raf);
  }, []);
  const el = scrollElRef.current;
  const sh = el?.scrollHeight ?? 0;
  const ch = el?.clientHeight ?? 0;
  const st = el ? Math.round(el.scrollTop) : 0;
  return (
    <div
      style={{
        position: 'fixed',
        top: 8,
        right: 8,
        zIndex: 99999,
        background: 'rgba(0,0,0,0.82)',
        color: '#19ff8a',
        font: '11px/1.4 monospace',
        padding: '6px 9px',
        borderRadius: 5,
        pointerEvents: 'none',
        whiteSpace: 'pre',
      }}
    >
      {`scrollHeight=${sh}  clientHeight=${ch}\nscrollTop=${st}\ntotalSize=${Math.round(getTotal())}\nSCROLLABLE=${sh > ch + 4 ? 'YES' : 'NO'}\natBottom=${atBottomRef.current}  revealed=${revealed}  ready=${ready}`}
    </div>
  );
}