import {
  forwardRef,
  useEffect,
  useImperativeHandle,
  useRef,
  useState,
} from 'react'
import { Box, Fab } from '@mui/material'
import { KeyboardArrowDown } from '@mui/icons-material'
import type { ChatContentProps, ChatContentRef } from '../types'
import { styles } from './styles'

/** Slack in px for treating the user as "at the bottom" when content grows. */
const AUTO_SCROLL_THRESHOLD = 32

export const ChatContent = forwardRef<ChatContentRef, ChatContentProps>(
  function ChatContent({ children, autoScroll = true, labels = {}, sx }, ref) {
    const scrollRef = useRef<HTMLDivElement>(null)
    const topSentinelRef = useRef<HTMLDivElement>(null)
    const bottomSentinelRef = useRef<HTMLDivElement>(null)

    const [isAtTop, setIsAtTop] = useState(true)
    const [isAtBottom, setIsAtBottom] = useState(true)

    const scrollToBottom = () => {
      scrollRef.current?.scrollTo({
        top: scrollRef.current.scrollHeight,
        behavior: 'smooth',
      })
    }

    const scrollToTop = () => {
      scrollRef.current?.scrollTo({ top: 0, behavior: 'smooth' })
    }

    useImperativeHandle(
      ref,
      () => ({ scrollToBottom, scrollToTop, isAtBottom, isAtTop }),
      [isAtBottom, isAtTop],
    )

    useEffect(() => {
      const root = scrollRef.current
      const topEl = topSentinelRef.current
      const bottomEl = bottomSentinelRef.current
      if (!root || !topEl || !bottomEl) return

      const topObserver = new IntersectionObserver(
        ([entry]) => setIsAtTop(entry?.isIntersecting ?? true),
        { root },
      )
      const bottomObserver = new IntersectionObserver(
        ([entry]) => setIsAtBottom(entry?.isIntersecting ?? true),
        { root },
      )
      topObserver.observe(topEl)
      bottomObserver.observe(bottomEl)
      return () => {
        topObserver.disconnect()
        bottomObserver.disconnect()
      }
    }, [])

    useEffect(() => {
      if (!autoScroll) return
      const root = scrollRef.current
      if (!root) return

      // We can't watch `children` to detect new content: tool traces and
      // `useTypewriter` live inside child components and update via local
      // state, so they mutate the DOM without re-rendering `ChatContent`.
      // MutationObserver picks up every DOM change regardless of source.
      let prevScrollHeight = root.scrollHeight
      let rafId: number | null = null

      const check = () => {
        rafId = null
        const newScrollHeight = root.scrollHeight
        if (newScrollHeight > prevScrollHeight) {
          const growth = newScrollHeight - prevScrollHeight
          // If the distance from the bottom doesn't exceed the amount the
          // content just grew (plus a small slack), the user was at the
          // bottom *before* the growth — scroll them to the new bottom.
          // Otherwise they've scrolled up to read history and we leave them
          // alone.
          const distanceFromBottom =
            newScrollHeight - root.scrollTop - root.clientHeight
          if (distanceFromBottom <= growth + AUTO_SCROLL_THRESHOLD) {
            root.scrollTo({ top: newScrollHeight, behavior: 'smooth' })
          }
        }
        prevScrollHeight = newScrollHeight
      }

      const mo = new MutationObserver(() => {
        // Plain `?? =` rather than `??=`: the React Compiler can't yet lower
        // logical-assignment operators, and bailing out would drop memoization
        // for this scroll container (which re-renders on every message).
        rafId = rafId ?? requestAnimationFrame(check)
      })
      mo.observe(root, {
        childList: true,
        subtree: true,
        characterData: true,
      })
      return () => {
        mo.disconnect()
        if (rafId !== null) cancelAnimationFrame(rafId)
      }
    }, [autoScroll])

    return (
      <Box
        ref={scrollRef}
        sx={{
          ...styles.content,
          borderTopColor: ({ palette }) =>
            isAtTop ? 'transparent' : palette.divider,
          borderBottomColor: ({ palette }) =>
            isAtBottom ? 'transparent' : palette.divider,
          ...sx,
        }}
      >
        <Box ref={topSentinelRef} sx={styles.sentinel} aria-hidden />
        {children}
        <Box ref={bottomSentinelRef} sx={styles.sentinel} aria-hidden />
        <Box sx={styles.jumpToLatestWrapper}>
          <Fab
            size='small'
            onClick={scrollToBottom}
            aria-label={labels.jumpToLatest ?? 'Jump to latest'}
            sx={{
              ...styles.jumpToLatest,
              opacity: isAtBottom ? 0 : 1,
              pointerEvents: isAtBottom ? 'none' : 'auto',
            }}
          >
            <KeyboardArrowDown />
          </Fab>
        </Box>
      </Box>
    )
  },
)
ChatContent.displayName = 'ChatContent'