src/components/ChatApp.tsx

import {
  useState,
  useRef,
  useEffect,
  useCallback,
  useLayoutEffect,
} from "react";
import {
  Send,
  Square,
  Plus,
  Lightbulb,
  LightbulbOff,
  ChevronDown,
} from "lucide-react";

import { useLLM } from "../hooks/useLLM";
import { MessageBubble } from "./MessageBubble";
import { MODEL_CONFIG } from "../model-config";
import type { ThinkingMode } from "../hooks/LLMContext";

const TEXTAREA_MIN_HEIGHT = "7.5rem";

const THINKING_OPTIONS: { value: ThinkingMode; label: string }[] = [
  { value: "enabled", label: "Reasoning On" },
  { value: "disabled", label: "Reasoning Off" },
];

export function ChatApp() {
  const {
    messages,
    isGenerating,
    tps,
    send,
    stop,
    status,
    clearChat,
    thinkingMode,
    setThinkingMode,
  } = useLLM();

  const [input, setInput] = useState("");
  const scrollRef = useRef<HTMLElement>(null);
  const textareaRef = useRef<HTMLTextAreaElement>(null);
  const [showThinkingMenu, setShowThinkingMenu] = useState(false);
  const thinkingMenuRef = useRef<HTMLDivElement>(null);
  const [thinkingSeconds, setThinkingSeconds] = useState(0);
  const thinkingStartRef = useRef<number | null>(null);
  const thinkingSecondsMapRef = useRef<Map<number, number>>(new Map());
  const prevIsGeneratingRef = useRef(false);
  const messagesRef = useRef(messages);
  const thinkingSecondsRef = useRef(thinkingSeconds);

  messagesRef.current = messages;
  thinkingSecondsRef.current = thinkingSeconds;

  const isReady = status.state === "ready";
  const hasMessages = messages.length > 0;

  const hasCompletedRef = useRef(false);
  useEffect(() => {
    if (hasMessages && !isGenerating) hasCompletedRef.current = true;
    if (!hasMessages) hasCompletedRef.current = false;
  }, [hasMessages, isGenerating]);

  const showNewChat =
    isReady && hasMessages && !isGenerating && hasCompletedRef.current;

  const prevMsgCountRef = useRef(0);
  const lastUserRef = useRef<HTMLDivElement>(null);
  const bottomSpacerRef = useRef<HTMLDivElement>(null);
  const userHasScrolledRef = useRef(false);

  useEffect(() => {
    const handleClickOutside = (event: MouseEvent) => {
      if (
        thinkingMenuRef.current &&
        !thinkingMenuRef.current.contains(event.target as Node)
      ) {
        setShowThinkingMenu(false);
      }
    };

    document.addEventListener("mousedown", handleClickOutside);
    return () => document.removeEventListener("mousedown", handleClickOutside);
  }, []);

  useEffect(() => {
    if (prevIsGeneratingRef.current && !isGenerating) {
      const lastMessage = messagesRef.current.at(-1);
      if (
        lastMessage?.role === "assistant" &&
        lastMessage.reasoning &&
        thinkingSecondsRef.current > 0
      ) {
        thinkingSecondsMapRef.current.set(
          lastMessage.id,
          thinkingSecondsRef.current,
        );
      }
    }
    prevIsGeneratingRef.current = isGenerating;
  }, [isGenerating]);

  useEffect(() => {
    if (!isGenerating) {
      thinkingStartRef.current = null;
      return;
    }

    thinkingStartRef.current = Date.now();
    setThinkingSeconds(0);

    const interval = setInterval(() => {
      if (thinkingStartRef.current) {
        setThinkingSeconds(
          Math.round((Date.now() - thinkingStartRef.current) / 1000),
        );
      }
    }, 500);

    return () => clearInterval(interval);
  }, [isGenerating]);

  const lastAssistant = messages.at(-1);
  useEffect(() => {
    if (
      isGenerating &&
      lastAssistant?.role === "assistant" &&
      lastAssistant.content
    ) {
      thinkingStartRef.current = null;
    }
  }, [isGenerating, lastAssistant?.role, lastAssistant?.content]);

  const getContainerPadTop = useCallback(() => {
    const container = scrollRef.current;
    if (!container) return 0;
    return parseFloat(getComputedStyle(container).paddingTop) || 0;
  }, []);

  const recalcSpacer = useCallback(() => {
    const container = scrollRef.current;
    const userElement = lastUserRef.current;
    const spacer = bottomSpacerRef.current;
    if (!container || !userElement || !spacer) return;

    // Compute the user element's offset from the top of the scroll content
    // (immune to browser scrollTop clamping).
    const userOffsetInContent =
      userElement.getBoundingClientRect().top -
      container.getBoundingClientRect().top +
      container.scrollTop;

    const padTop = getContainerPadTop();
    const padBottom =
      parseFloat(getComputedStyle(container).paddingBottom) || 0;
    const usableHeight = container.clientHeight - padTop - padBottom;

    // Distance from user element top to the spacer top (everything between them).
    const contentBelowUser =
      spacer.getBoundingClientRect().top -
      userElement.getBoundingClientRect().top;

    spacer.style.height = `${Math.max(0, usableHeight - contentBelowUser)}px`;

    // Only re-anchor if the user hasn't manually scrolled away.
    if (!userHasScrolledRef.current) {
      const desiredScrollTop = userOffsetInContent - padTop;
      if (Math.abs(container.scrollTop - desiredScrollTop) > 0.5) {
        container.scrollTo({ top: desiredScrollTop, behavior: "smooth" });
      }
    }
  }, [getContainerPadTop]);

  useLayoutEffect(() => {
    recalcSpacer();

    const isNewMessage = messages.length > prevMsgCountRef.current;
    prevMsgCountRef.current = messages.length;

    if (isNewMessage) {
      // Reset manual-scroll flag so the initial snap works for the new message.
      userHasScrolledRef.current = false;

      const container = scrollRef.current;
      const userElement = lastUserRef.current;
      if (!container || !userElement) return;

      const scrollTarget =
        container.scrollTop +
        (userElement.getBoundingClientRect().top -
          container.getBoundingClientRect().top) -
        getContainerPadTop();
      container.scrollTo({ top: scrollTarget, behavior: "smooth" });
    }
  }, [messages, isGenerating, recalcSpacer, getContainerPadTop]);

  useEffect(() => {
    window.addEventListener("resize", recalcSpacer);
    return () => window.removeEventListener("resize", recalcSpacer);
  }, [recalcSpacer]);

  // Detect user-initiated scrolls (wheel or touch) to unlock free scrolling
  // during generation.
  useEffect(() => {
    const container = scrollRef.current;
    if (!container) return;

    const markScrolled = () => {
      if (isGenerating) {
        userHasScrolledRef.current = true;
      }
    };

    container.addEventListener("wheel", markScrolled, { passive: true });
    container.addEventListener("touchmove", markScrolled, { passive: true });
    return () => {
      container.removeEventListener("wheel", markScrolled);
      container.removeEventListener("touchmove", markScrolled);
    };
  }, [isGenerating]);

  // Re-anchor the user message whenever the scroll container's own size changes
  useLayoutEffect(() => {
    const container = scrollRef.current;
    if (!container) return;

    let lastHeight = container.clientHeight;
    const observer = new ResizeObserver(() => {
      const h = container.clientHeight;
      if (h !== lastHeight) {
        lastHeight = h;
        recalcSpacer();
      }
    });
    observer.observe(container);
    return () => observer.disconnect();
  }, [recalcSpacer]);

  const handleSubmit = useCallback(
    (event?: React.FormEvent) => {
      event?.preventDefault();
      const text = input.trim();
      if (!text || !isReady || isGenerating) return;
      setInput("");
      if (textareaRef.current) {
        textareaRef.current.style.height = TEXTAREA_MIN_HEIGHT;
      }
      send(text);
    },
    [input, isReady, isGenerating, send],
  );

  const handleInputKeyDown = useCallback(
    (event: React.KeyboardEvent<HTMLTextAreaElement>) => {
      if (event.key === "Enter" && !event.shiftKey) {
        event.preventDefault();
        handleSubmit();
      }
    },
    [handleSubmit],
  );

  const lastUserIndex = messages.findLastIndex(
    (message) => message.role === "user",
  );

  const renderInputArea = (showDisclaimer: boolean) => (
    <div className="w-full">
      <form onSubmit={handleSubmit} className="mx-auto max-w-3xl">
        <div className="relative">
          <textarea
            ref={textareaRef}
            className="w-full rounded-xl border border-[#2a2a2a] bg-[#111111] px-4 py-3 pb-11 text-[15px] text-[#f0f0f0] placeholder-[#888888] focus:border-[#c9a84c] focus:outline-none focus:ring-1 focus:ring-[#c9a84c] disabled:opacity-50 resize-none max-h-40 shadow-sm"
            style={{
              minHeight: TEXTAREA_MIN_HEIGHT,
              height: TEXTAREA_MIN_HEIGHT,
            }}
            placeholder={isReady ? "Type a message…" : "Loading model…"}
            value={input}
            onChange={(event) => {
              setInput(event.target.value);
              event.target.style.height = TEXTAREA_MIN_HEIGHT;
              event.target.style.height =
                Math.max(event.target.scrollHeight, 120) + "px";
            }}
            onKeyDown={handleInputKeyDown}
            disabled={!isReady}
            autoFocus
          />

          <div className="absolute bottom-2 left-2 right-2 flex items-center justify-between pb-3 px-2">
            <div ref={thinkingMenuRef} className="relative">
              <button
                type="button"
                onClick={() => setShowThinkingMenu((value) => !value)}
                className="flex items-center gap-1 rounded-lg text-xs text-[#888888] hover:text-[#f0f0f0] transition-colors cursor-pointer"
                title="Reasoning mode"
              >
                {thinkingMode === "enabled" ? (
                  <Lightbulb className="h-3.5 w-3.5 text-[#76b900]" />
                ) : (
                  <LightbulbOff className="h-3.5 w-3.5 text-[#888888]" />
                )}
                <span>
                  {
                    THINKING_OPTIONS.find(
                      (option) => option.value === thinkingMode,
                    )?.label
                  }
                </span>
                <ChevronDown className="h-3 w-3" />
              </button>

              {showThinkingMenu && (
                <div className="absolute bottom-full left-0 mb-1 min-w-[160px] rounded-lg border border-[#2a2a2a] bg-[#1a1a1a] py-1 shadow-xl z-50">
                  {THINKING_OPTIONS.map((option) => (
                    <button
                      key={option.value}
                      type="button"
                      onClick={() => {
                        setThinkingMode(option.value);
                        setShowThinkingMenu(false);
                      }}
                      className={`flex w-full items-center gap-2 px-3 py-1.5 text-left text-xs transition-colors cursor-pointer ${
                        thinkingMode === option.value
                          ? "bg-[#222222] text-[#f0f0f0]"
                          : "text-[#888888] hover:bg-[#222222] hover:text-[#f0f0f0]"
                      }`}
                    >
                      {option.value === "enabled" ? (
                        <Lightbulb className="h-3.5 w-3.5 text-[#76b900]" />
                      ) : (
                        <LightbulbOff className="h-3.5 w-3.5 text-[#888888]" />
                      )}
                      <span>{option.label}</span>
                    </button>
                  ))}
                </div>
              )}
            </div>

            {isGenerating ? (
              <button
                type="button"
                onClick={stop}
                className="flex items-center justify-center rounded-lg text-[#76b900] hover:text-[#8fd400] transition-colors cursor-pointer"
                title="Stop generating"
              >
                <Square className="h-4 w-4 fill-current" />
              </button>
            ) : (
              <button
                type="submit"
                disabled={!isReady || !input.trim()}
                className="flex items-center justify-center rounded-lg text-[#76b900] hover:text-[#8fd400] disabled:opacity-30 transition-colors cursor-pointer"
                title="Send message"
              >
                <Send className="h-4 w-4" />
              </button>
            )}
          </div>
        </div>
      </form>

      {showDisclaimer && (
        <p className="mx-auto max-w-3xl mt-1 text-center text-xs text-[#888888]">
          No chats are sent to a server. Everything runs locally in your
          browser. AI can make mistakes. Check important info.
        </p>
      )}
    </div>
  );

  return (
    <div className="flex h-full flex-col bg-[#0a0a0a] text-[#f0f0f0]">
      <header className="flex-none flex items-center justify-between border-b border-[#2a2a2a] px-6 py-3 h-14">
        <div className="flex items-center gap-2">
          <svg
            height="16"
            viewBox="0 0 20 20"
            width="16"
            fill="none"
            xmlns="http://www.w3.org/2000/svg"
            aria-hidden="true"
          >
            <path d="M10 1.5 2 6v8l8 4.5 8-4.5V6l-8-4.5Z" fill="#76B900" />
            <path d="M10 5.1 6 7.4v5.2l4 2.3 4-2.3V7.4l-4-2.3Z" fill="#fff" />
          </svg>
          <h1 className="text-base font-semibold text-[#f0f0f0]">
            {MODEL_CONFIG.label}
          </h1>
          <span className="text-base font-semibold text-[#888888]">WebGPU</span>
        </div>
        <div className="flex items-center gap-3">
          <button
            onClick={clearChat}
            className={`flex items-center gap-1.5 rounded-lg px-2.5 py-1.5 text-xs text-[#888888] hover:text-[#f0f0f0] hover:bg-[#1a1a1a] transition-opacity duration-300 cursor-pointer ${
              showNewChat ? "opacity-100" : "opacity-0 pointer-events-none"
            }`}
            title="New chat"
          >
            <Plus className="h-3.5 w-3.5" />
            New chat
          </button>
        </div>
      </header>

      {isReady && !hasMessages ? (
        <div className="flex flex-1 flex-col items-center justify-center px-4">
          <div className="mb-8 text-center">
            <p className="text-3xl font-medium text-[#f0f0f0]">
              What can I help you with?
            </p>
          </div>

          {renderInputArea(false)}

          <div className="mt-6 flex flex-wrap justify-center gap-3 max-w-3xl">
            {MODEL_CONFIG.examplePrompts.map(({ label, prompt }) => (
              <button
                key={label}
                onClick={() => send(prompt)}
                className="relative rounded-[14px] border-[2px] border-[rgba(220,210,180,0.3)] bg-[linear-gradient(160deg,rgba(28,28,28,0.97)_0%,rgba(10,10,10,0.99)_100%)] px-4 py-2.5 text-xs text-[#888888] hover:text-[#f0f0f0] hover:border-[rgba(235,225,195,0.6)] hover:-translate-y-[2px] hover:shadow-[0_4px_16px_rgba(0,0,0,0.4),0_0_10px_rgba(220,210,180,0.05)] transition-all duration-300 cursor-pointer overflow-hidden before:pointer-events-none before:absolute before:inset-0 before:rounded-[14px] before:bg-[radial-gradient(ellipse_at_50%_0%,rgba(235,225,190,0.05)_0%,transparent_55%)]"
              >
                {label}
              </button>
            ))}
          </div>
        </div>
      ) : (
        <>
          <main
            ref={scrollRef}
            className="min-h-0 flex-1 overflow-y-auto px-4 py-6 animate-fade-in"
          >
            <div className="mx-auto flex max-w-3xl flex-col gap-6">
              {messages.map((message, index) => {
                const isLastAssistant =
                  index === messages.length - 1 && message.role === "assistant";
                const isLastUser =
                  message.role === "user" && index === lastUserIndex;

                return (
                  <div
                    key={message.id}
                    ref={isLastUser ? lastUserRef : undefined}
                  >
                    <MessageBubble
                      msg={message}
                      index={index}
                      isStreaming={isGenerating && isLastAssistant}
                      thinkingSeconds={
                        isLastAssistant
                          ? thinkingSeconds
                          : thinkingSecondsMapRef.current.get(message.id)
                      }
                      isGenerating={isGenerating}
                    />
                  </div>
                );
              })}
              <div ref={bottomSpacerRef} />
            </div>
          </main>

          <footer className="flex-none px-4 py-3 animate-fade-in">
            <p
              className={`mx-auto max-w-3xl mb-3 text-center text-xs font-mono tabular-nums h-4 ${isGenerating && tps > 0 ? "text-[#888888]" : "text-transparent"}`}
            >
              {isGenerating && tps > 0
                ? `${tps.toFixed(1)} tokens/s`
                : "\u00A0"}
            </p>
            {renderInputArea(true)}
          </footer>
        </>
      )}
    </div>
  );
}