import { useLayoutEffect, useMemo, useRef, useState } from 'react';
import { ArrowDown, Loader2 } from 'lucide-react';
import { useVirtualizer } from '@tanstack/react-virtual';
import { measureHeight } from 'helpers/measure-zoomed';
import type { UseEmbeddableChatType } from '../EmbeddableChat/useEmbeddableChat';
import { MessageBubble, StreamingBubble } from './MessageBubble';
import { buildRows } from './turn-rows';
import { useFoldExit } from './fold-exit';
import { ToolOpenProvider } from './tool-open';
import { TurnEnd, TurnFold } from './TurnFold';

const OVERSCAN = 5;

type MessageListProps = {
  manager: UseEmbeddableChatType;
};

export const MessageList = ({ manager }: MessageListProps) => {
  const {
    messages,
    historicalCount,
    streamingText,
    isGenerating,
    showJumpToBottom,
    jumpToBottom,
    sendPrompt,
    scrollViewportRef,
    bottomRef,
    isLoadingOlder,
    stickToBottom,
  } = manager;

  // Turns from this sitting stay open; anything loaded from the server folds. See `turn-rows.ts`.
  const rows = useMemo(
    () => buildRows(messages, isGenerating, historicalCount),
    [messages, isGenerating, historicalCount],
  );

  // Which folds you have opened. Lifted out of the fold itself because rows are virtualised: state held
  // inside one would be thrown away the moment it scrolled past the overscan window and silently snap shut.
  const [openFolds, setOpenFolds] = useState<ReadonlySet<string>>(() => new Set());
  const toggleFold = (key: string) =>
    setOpenFolds((prev) => {
      const next = new Set(prev);
      if (!next.delete(key)) next.add(key);
      return next;
    });

  // Whether a turn has run while this list has been mounted. Tool rows expand themselves so you can watch
  // work happen, which is only a service to someone who is watching — reopening yesterday's conversation
  // should hand you a readable page, not every output it ever produced. Never reset: sending the next
  // message starts generating again, so the flag is already true by the time the new turn has any rows.
  const watched = useRef(false);
  if (isGenerating) watched.current = true;

  // The turn you just superseded, on its way out. It holds the work open for a beat and then shrinks it.
  const foldExit = useFoldExit(rows, watched.current);

  const virtualizer = useVirtualizer({
    count: rows.length,
    getScrollElement: () => scrollViewportRef.current,
    estimateSize: () => 150,
    overscan: OVERSCAN,
    // Rows must reconcile by identity, not by position. The default extractor keys them by index, so
    // prepending older messages hands every row the key of the row above it — React then keeps the mounted
    // component and swaps its props, and the measured height cache follows the slot instead of the content.
    getItemKey: (index) => rows[index]?.key ?? index,
    // Measured height must be in the same unit as the translateY below — layout px, which is what the
    // panel zoom leaves alone. A raw getBoundingClientRect is in rendered px and stacks the bubbles.
    measureElement: measureHeight,
  });

  // Re-pin once the rows have actually measured themselves.
  //
  // A new row is laid out at the 150px estimate, so the scroll that fires when the message arrives lands on
  // a bottom that isn't the real one — and the correction arrives later, through a ResizeObserver, on a
  // render that no message change accompanies. Nothing was watching for it. An expanded tool call is several
  // times the estimate, which is why this got worse the moment tool rows stopped collapsing on a timer:
  // every row of a running turn now overshoots, and the view sat a screenful short of the work.
  const totalSize = virtualizer.getTotalSize();
  useLayoutEffect(() => {
    stickToBottom();
  }, [totalSize, stickToBottom]);

  // Keep the viewport anchored when older messages are prepended: the row that used to be first (now
  // shifted down by however many rows arrived above it) is pinned back to the top, so content doesn't jump.
  // Counted in rows rather than messages, because a prepended turn folds down to fewer rows than it has
  // messages and scrolling by the message count would overshoot.
  const prevLenRef = useRef(rows.length);
  const prevFirstKeyRef = useRef<string | null>(null);
  useLayoutEffect(() => {
    const grew = rows.length - prevLenRef.current;
    const firstKey = rows[0]?.key ?? '';
    const prepended = grew > 0 && prevFirstKeyRef.current !== null && firstKey !== prevFirstKeyRef.current;
    if (prepended) virtualizer.scrollToIndex(grew, { align: 'start' });
    prevLenRef.current = rows.length;
    prevFirstKeyRef.current = firstKey;
  }, [rows, virtualizer]);

  return (
    <div className="flex-1 min-h-0 relative">
      {isLoadingOlder && (
        <div className="absolute top-1 left-1/2 -translate-x-1/2 z-10 rounded-full bg-background/90 shadow p-1.5">
          <Loader2 className="h-4 w-4 animate-spin text-duck-teal" />
        </div>
      )}
      <div ref={scrollViewportRef} className="h-full overflow-y-auto">
        {rows.length === 0 && !isGenerating && (
          <div className="flex items-center justify-center h-full min-h-[200px] text-muted-foreground text-sm">
            Send a message to start
          </div>
        )}
        {rows.length > 0 && (
          <ToolOpenProvider>
            <div className="relative w-full" style={{ height: virtualizer.getTotalSize() }}>
              {virtualizer.getVirtualItems().map((virtualRow) => {
                const row = rows[virtualRow.index]!;
                return (
                  <div
                    key={virtualRow.key}
                    data-index={virtualRow.index}
                    ref={virtualizer.measureElement}
                    className="absolute top-0 left-0 w-full px-4"
                    style={{
                      transform: `translateY(${virtualRow.start}px)`,
                    }}
                  >
                    <div className="py-1.5">
                      {row.kind === 'message' && (
                        <MessageBubble
                          message={row.message}
                          onAnswer={(text) => sendPrompt(text)}
                          onRetry={(prompt) => sendPrompt(prompt)}
                          // Deliberately not `isGenerating`: a turn's work stays open after it finishes, and
                          // closes when you send the next message and the turn folds behind its summary.
                          defaultOpen={row.live && watched.current}
                        />
                      )}
                      {row.kind === 'fold' && (
                        <TurnFold
                          messages={row.messages}
                          counts={row.counts}
                          // While a fold is on its way out it answers to the animation rather than to you;
                          // your own toggles take back over the moment it has finished closing.
                          open={foldExit?.key === row.key ? foldExit.open : openFolds.has(row.key)}
                          childrenOpen={foldExit?.key === row.key}
                          onToggle={() => toggleFold(row.key)}
                        />
                      )}
                      {row.kind === 'turn-end' && <TurnEnd />}
                    </div>
                  </div>
                );
              })}
            </div>
          </ToolOpenProvider>
        )}
        {isGenerating && (
          <div className="px-4 py-1.5">
            <StreamingBubble text={streamingText} />
          </div>
        )}
        <div ref={bottomRef} />
      </div>

      {showJumpToBottom && (
        <button
          onClick={jumpToBottom}
          aria-label="Jump to latest message"
          className="absolute bottom-2 left-1/2 -translate-x-1/2 cursor-pointer rounded-full bg-primary p-1.5 text-primary-foreground shadow-lg transition-colors hover:bg-primary/90"
        >
          <ArrowDown className="h-4 w-4" />
        </button>
      )}
    </div>
  );
};
