src/components/ChatApp.tsx

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

import { useLLM } from "../hooks/useLLM";
import { MessageBubble } from "./MessageBubble";
import { BrandMark } from "./BrandMark";
import { MODEL_CONFIG } from "../model-config";

const TEXTAREA_MIN_HEIGHT = "7.5rem";

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

  const [input, setInput] = useState("");
  const scrollRef = useRef<HTMLElement>(null);
  const textareaRef = useRef<HTMLTextAreaElement>(null);

  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);

  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;

    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;

    const contentBelowUser = spacer.getBoundingClientRect().top - userElement.getBoundingClientRect().top;

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

    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) {
      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]);

  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]);

  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 max-w-3xl mx-auto">
      <form onSubmit={handleSubmit} className="relative">
        <textarea
          ref={textareaRef}
          className="w-full min-h-[7.5rem] pt-3.5 px-4 pb-11 border border-line rounded-[20px] bg-[rgba(255,255,255,0.04)] text-text font-body text-[0.95rem] leading-[1.5] resize-none max-h-40 placeholder:text-text-muted focus:outline-[1px] focus:outline-[rgba(157,224,255,0.44)] focus:border-[rgba(157,224,255,0.44)] disabled:opacity-50"
          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-end px-2 pb-3">
          {isGenerating ? (
            <button
              type="button"
              onClick={stop}
              className="flex items-center justify-center p-1.5 rounded-lg text-accent bg-transparent transition-[color,opacity] duration-[180ms] ease-[ease] hover:text-accent-strong"
              title="Stop generating"
            >
              <Square size={16} fill="currentColor" />
            </button>
          ) : (
            <button
              type="submit"
              disabled={!isReady || !input.trim()}
              className="flex items-center justify-center p-1.5 rounded-lg text-accent bg-transparent transition-[color,opacity] duration-[180ms] ease-[ease] hover:text-accent-strong disabled:opacity-30"
              title="Send message"
            >
              <Send size={16} />
            </button>
          )}
        </div>
      </form>

      {showDisclaimer && (
        <p className="max-w-3xl mx-auto mt-1 text-center text-[0.75rem] text-text-muted">
          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 flex-col h-full bg-bg text-text">
      <header className="flex-none flex items-center justify-between gap-4 py-4 px-6 border-b border-line max-sm:py-3 max-sm:px-4">
        <BrandMark />
        <button
          onClick={clearChat}
          className={`inline-flex items-center gap-2.5 py-2.5 px-4 border border-line rounded-full bg-[rgba(8,13,24,0.44)] text-text text-[0.88rem] backdrop-blur-[16px] transition-[transform,border-color,background-color,opacity] duration-[180ms] ease-[ease] hover:-translate-y-0.5 ${
            showNewChat ? "" : "opacity-0 pointer-events-none"
          }`}
          title="New chat"
        >
          <Plus className="shrink-0 text-accent" size={16} strokeWidth={1.8} />
          New chat
        </button>
      </header>

      {isReady && !hasMessages ? (
        <div className="flex-1 flex flex-col items-center justify-center p-4">
          <h2 className="mb-8 text-[clamp(1.6rem,4vw,2.4rem)] font-bold tracking-[-0.04em] text-center">
            What can I help you with?
          </h2>

          {renderInputArea(false)}

          <div className="flex flex-wrap justify-center gap-2.5 max-w-3xl mt-6">
            {MODEL_CONFIG.examplePrompts.map(({ label, prompt }) => (
              <button
                key={label}
                onClick={() => send(prompt)}
                className="py-2.5 px-3.5 border border-line rounded-full bg-[rgba(255,255,255,0.05)] text-text-soft text-[0.88rem] transition-[transform,border-color,background-color,opacity] duration-[180ms] ease-[ease] hover:-translate-y-0.5 hover:border-[rgba(157,224,255,0.42)] hover:bg-[rgba(157,224,255,0.14)] hover:text-accent-strong"
                type="button"
              >
                {label}
              </button>
            ))}
          </div>
        </div>
      ) : (
        <>
          <main ref={scrollRef} className="flex-1 min-h-0 overflow-y-auto py-6 px-4 animate-fade-in">
            <div className="max-w-3xl mx-auto flex 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}
                      isGenerating={isGenerating}
                    />
                  </div>
                );
              })}
              <div ref={bottomSpacerRef} />
            </div>
          </main>

          <footer className="flex-none pt-3 px-4 pb-4 animate-fade-in">
            {isGenerating && tps > 0 && (
              <p className="max-w-3xl mx-auto mb-3 text-center font-mono text-[0.78rem] tabular-nums h-4 text-text-muted">
                {`${tps.toFixed(1)} tokens/s`}
              </p>
            )}
            {renderInputArea(true)}
          </footer>
        </>
      )}
    </div>
  );
}