'use client'

import { useEffect, useMemo, useRef, useState } from 'react'

import { UseChatHelpers } from '@ai-sdk/react'

import { useMediaQuery } from '@/lib/hooks/use-media-query'
import type { UIDataTypes, UIMessage, UITools } from '@/lib/types/ai'
import { cn } from '@/lib/utils'
import { extractCitationMapsFromMessages } from '@/lib/utils/citation'

import { AnimatedLogo } from './ui/animated-logo'
import { ChatError } from './chat-error'
import { ChatFooterMessage } from './chat-footer-message'
import { RenderMessage } from './render-message'

// Import section structure interface
interface ChatSection {
  id: string
  userMessage: UIMessage
  assistantMessages: UIMessage[]
}

interface ChatMessagesProps {
  sections: ChatSection[] // Changed from messages to sections
  status: UseChatHelpers<UIMessage<unknown, UIDataTypes, UITools>>['status']
  chatId?: string
  isGuest?: boolean
  isCloudDeployment?: boolean
  libraryAvailable?: boolean
  addToolResult?: (params: { toolCallId: string; result: any }) => void
  /** Ref for the scroll container */
  scrollContainerRef: React.RefObject<HTMLDivElement>
  onUpdateMessage?: (messageId: string, newContent: string) => Promise<void>
  reload?: (messageId: string) => Promise<void | string | null | undefined>
  error?: Error | string | null | undefined
  onQuoteContext?: (text: string) => void
}

const DESKTOP_LATEST_SECTION_OFFSET = 196
const MOBILE_LATEST_SECTION_OFFSET_FALLBACK = 180
const MOBILE_FOLLOW_UP_TOP_CLEARANCE_FALLBACK = 56

export function ChatMessages({
  sections,
  status,
  chatId,
  isGuest = false,
  isCloudDeployment = false,
  libraryAvailable = true,
  addToolResult,
  scrollContainerRef,
  onUpdateMessage,
  reload,
  error,
  onQuoteContext
}: ChatMessagesProps) {
  // Track user-modified states (when user explicitly opens/closes)
  const [userModifiedStates, setUserModifiedStates] = useState<
    Record<string, boolean>
  >({})
  // Cache tool counts for performance optimization
  const toolCountCacheRef = useRef<Map<string, number>>(new Map())
  const isLoading = status === 'submitted' || status === 'streaming'
  const isMobile = useMediaQuery('(max-width: 767px)')
  const [scrollViewportHeight, setScrollViewportHeight] = useState(0)
  const [mobileFollowUpTopClearance, setMobileFollowUpTopClearance] = useState(
    MOBILE_FOLLOW_UP_TOP_CLEARANCE_FALLBACK
  )

  // Tool types definition - moved outside function for performance
  const toolTypes = ['tool-search', 'tool-fetch', 'tool-askQuestion']

  useEffect(() => {
    const container = scrollContainerRef.current
    if (!container) return

    const updateMetrics = () => {
      const { paddingTop } = getComputedStyle(container)
      setScrollViewportHeight(container.clientHeight)
      setMobileFollowUpTopClearance(
        parseFloat(paddingTop) || MOBILE_FOLLOW_UP_TOP_CLEARANCE_FALLBACK
      )
    }
    updateMetrics()

    if (typeof ResizeObserver === 'undefined') return

    const observer = new ResizeObserver(updateMetrics)
    observer.observe(container)

    return () => observer.disconnect()
  }, [scrollContainerRef])

  // Clear cache during streaming to ensure accurate tool counts
  useEffect(() => {
    if (isLoading) {
      // Clear cache for all messages during streaming
      toolCountCacheRef.current.clear()
    }
  }, [isLoading])

  const latestSectionMinHeight =
    isMobile && scrollViewportHeight > 0
      ? `${Math.max(0, scrollViewportHeight - mobileFollowUpTopClearance)}px`
      : `calc(100dvh - ${
          isMobile
            ? MOBILE_LATEST_SECTION_OFFSET_FALLBACK
            : DESKTOP_LATEST_SECTION_OFFSET
        }px)`

  // Extract citation maps from all messages in all sections
  const allCitationMaps = useMemo(() => {
    const allMessages: UIMessage[] = []
    sections.forEach(section => {
      allMessages.push(section.userMessage)
      allMessages.push(...section.assistantMessages)
    })
    return extractCitationMapsFromMessages(allMessages)
  }, [sections])

  if (!sections.length) return null

  // Keep the assistant logo visible for the latest section after generation
  const latestSection = sections.at(-1)
  const showAssistantLogo = Boolean(
    latestSection && (isLoading || latestSection.assistantMessages.length > 0)
  )

  // Helper function to get tool count with caching
  const getToolCount = (message?: UIMessage): number => {
    if (!message || !message.id) return 0

    // During streaming, always recalculate
    if (isLoading) {
      const count =
        message.parts?.filter(part => toolTypes.includes(part.type)).length || 0
      return count
    }

    // Check cache first when not streaming
    const cached = toolCountCacheRef.current.get(message.id)
    if (cached !== undefined) {
      return cached
    }

    // Calculate and cache
    const count =
      message.parts?.filter(part => toolTypes.includes(part.type)).length || 0
    toolCountCacheRef.current.set(message.id, count)
    return count
  }

  const getIsOpen = (
    id: string,
    partType?: string,
    hasNextPart?: boolean,
    message?: UIMessage
  ) => {
    // If user has explicitly modified this state, use that
    if (userModifiedStates.hasOwnProperty(id)) {
      return userModifiedStates[id]
    }

    // For tool types, check if there are multiple tools
    if (partType && toolTypes.includes(partType)) {
      const toolCount = getToolCount(message)
      // If multiple tools exist, default to closed
      if (toolCount > 1) {
        return false
      }
      // Keep a single tool open until answer text follows it
      return !hasNextPart
    }

    // For tool-invocations, default to open
    if (partType === 'tool-invocation') {
      return true
    }

    // For reasoning, auto-collapse if there's a next part in the same message
    if (partType === 'reasoning') {
      return !hasNextPart
    }

    // For other types (like text), default to open
    return true
  }

  const handleOpenChange = (id: string, open: boolean) => {
    setUserModifiedStates(prev => ({
      ...prev,
      [id]: open
    }))
  }

  return (
    <div
      id="scroll-container"
      ref={scrollContainerRef}
      role="list"
      aria-roledescription="chat messages"
      className={cn(
        'relative size-full pt-14',
        sections.length > 0 ? 'flex-1 overflow-y-auto' : ''
      )}
    >
      <div className="relative mx-auto w-full max-w-full md:max-w-3xl px-4">
        {sections.map((section, sectionIndex) => (
          <div
            key={section.id}
            id={`section-${section.id}`}
            className="chat-section scroll-mt-14 pb-4 md:pb-14"
            style={
              sectionIndex === sections.length - 1
                ? { minHeight: latestSectionMinHeight }
                : {}
            }
          >
            {/* User message */}
            <div className="flex flex-col gap-2 md:gap-4 mb-2 md:mb-4">
              <RenderMessage
                message={section.userMessage}
                messageId={section.userMessage.id}
                getIsOpen={(id, partType, hasNextPart) =>
                  getIsOpen(id, partType, hasNextPart, section.userMessage)
                }
                onOpenChange={handleOpenChange}
                chatId={chatId}
                isGuest={isGuest}
                isCloudDeployment={isCloudDeployment}
                libraryAvailable={libraryAvailable}
                status={status}
                addToolResult={addToolResult}
                onUpdateMessage={onUpdateMessage}
                reload={reload}
                citationMaps={allCitationMaps}
                onQuoteContext={onQuoteContext}
              />
            </div>

            {/* Assistant messages */}
            {section.assistantMessages.map((assistantMessage, messageIndex) => {
              // Check if this is the latest assistant message in the latest section
              const isLatestMessage =
                sectionIndex === sections.length - 1 &&
                messageIndex === section.assistantMessages.length - 1

              return (
                <div
                  key={assistantMessage.id}
                  className="flex flex-col gap-2 md:gap-4"
                >
                  <RenderMessage
                    message={assistantMessage}
                    messageId={assistantMessage.id}
                    getIsOpen={(id, partType, hasNextPart) =>
                      getIsOpen(id, partType, hasNextPart, assistantMessage)
                    }
                    onOpenChange={handleOpenChange}
                    chatId={chatId}
                    isGuest={isGuest}
                    isCloudDeployment={isCloudDeployment}
                    libraryAvailable={libraryAvailable}
                    status={status}
                    addToolResult={addToolResult}
                    onUpdateMessage={onUpdateMessage}
                    reload={reload}
                    isLatestMessage={isLatestMessage}
                    citationMaps={allCitationMaps}
                    onQuoteContext={onQuoteContext}
                  />
                </div>
              )
            })}
            {/* Show assistant logo and footer message after assistant messages */}
            {showAssistantLogo && sectionIndex === sections.length - 1 && (
              <div className="flex items-center gap-3 py-1 md:py-4">
                <AnimatedLogo
                  className="size-10 shrink-0"
                  animate={isLoading}
                />
                <ChatFooterMessage isLoading={isLoading} />
              </div>
            )}
            {sectionIndex === sections.length - 1 && (
              <ChatError error={error} />
            )}
          </div>
        ))}
      </div>
    </div>
  )
}