import type { UIMessage } from "ai";

import { ArrowDownOutlined } from "@ant-design/icons";
import { Button } from "antd";
import "overlayscrollbars/styles/overlayscrollbars.css";
import { OverlayScrollbarsComponent, type OverlayScrollbarsComponentRef } from "overlayscrollbars-react";
import { useCallback, useRef, useState } from "react";

import { useChatScroll } from "./use-chat-scroll";

interface ChatScrollAreaProps {
  children: React.ReactNode;
  messages: UIMessage[];
  status: string;
}

export function ChatScrollArea({ children, messages, status }: ChatScrollAreaProps) {
  const scrollRef = useRef<HTMLDivElement>(null);
  const osRef = useRef<OverlayScrollbarsComponentRef>(null);
  const [viewportElement, setViewportElement] = useState<HTMLDivElement | null>(null);

  const handleOsInitialized = useCallback(() => {
    const os = osRef.current;
    if (!os) return;
    const instance = os.osInstance();
    if (!instance) return;
    const viewport = instance.elements().viewport as HTMLDivElement;
    scrollRef.current = viewport;
    setViewportElement(viewport);
  }, []);

  const { isAtBottom, scrollToBottom } = useChatScroll({ messages, scrollRef, status, viewportElement });

  return (
    <>
      <OverlayScrollbarsComponent
        className="chat-scroll-area"
        events={{ initialized: handleOsInitialized }}
        options={{
          overflow: { x: "hidden", y: "scroll" },
          scrollbars: { autoHide: "move", theme: "os-theme-custom" },
        }}
        ref={osRef}
      >
        {children}
      </OverlayScrollbarsComponent>
      {!isAtBottom && (
        <Button
          className="chat-scroll-bottom-btn"
          icon={<ArrowDownOutlined />}
          onClick={scrollToBottom}
          shape="circle"
          size="small"
        />
      )}
    </>
  );
}