src/App.jsx

import { useEffect, useState, useRef } from "react";

import Chat from "./components/Chat";
import ArrowRightIcon from "./components/icons/ArrowRightIcon";
import StopIcon from "./components/icons/StopIcon";
import Progress from "./components/Progress";

const IS_WEBGPU_AVAILABLE = !!navigator.gpu;
const STICKY_SCROLL_THRESHOLD = 120;
const EXAMPLES = [
  "Find all complex solutions of z^3 + 2z^2 - z - 2 = 0",
  "A train travels at 60 km/h for the first half of its journey and at 40 km/h for the second half. If the total journey is 300 km, how long does it take?",
  "Write a Python function to check if a given string is a palindrome, ignoring spaces and punctuation.",
];

// Add gradient animation classes
const GRADIENT_CLASSES = "animate-gradient bg-gradient-to-r from-blue-500 via-purple-500 to-pink-500 bg-[length:200%_auto]";
const CARD_HOVER_CLASSES = "hover:scale-105 hover:shadow-xl hover:shadow-blue-500/20 dark:hover:shadow-blue-800/20";

function App() {
  // Create a reference to the worker object.
  const worker = useRef(null);

  const textareaRef = useRef(null);
  const chatContainerRef = useRef(null);

  // Model loading and progress
  const [status, setStatus] = useState(null);
  const [error, setError] = useState(null);
  const [loadingMessage, setLoadingMessage] = useState("");
  const [progressItems, setProgressItems] = useState([]);
  const [isRunning, setIsRunning] = useState(false);

  // Inputs and outputs
  const [input, setInput] = useState("");
  const [messages, setMessages] = useState([]);
  const [tps, setTps] = useState(null);
  const [numTokens, setNumTokens] = useState(null);

  function onEnter(message) {
    setMessages((prev) => [...prev, { role: "user", content: message }]);
    setTps(null);
    setIsRunning(true);
    setInput("");
  }

  function onInterrupt() {
    // NOTE: We do not set isRunning to false here because the worker
    // will send a 'complete' message when it is done.
    worker.current.postMessage({ type: "interrupt" });
  }

  useEffect(() => {
    resizeInput();
  }, [input]);

  function resizeInput() {
    if (!textareaRef.current) return;

    const target = textareaRef.current;
    target.style.height = "auto";
    const newHeight = Math.min(Math.max(target.scrollHeight, 24), 200);
    target.style.height = `${newHeight}px`;
  }

  // We use the `useEffect` hook to setup the worker as soon as the `App` component is mounted.
  useEffect(() => {
    // Create the worker if it does not yet exist.
    if (!worker.current) {
      worker.current = new Worker(new URL("./worker.js", import.meta.url), {
        type: "module",
      });
      worker.current.postMessage({ type: "check" }); // Do a feature check
    }

    // Create a callback function for messages from the worker thread.
    const onMessageReceived = (e) => {
      switch (e.data.status) {
        case "loading":
          // Model file start load: add a new progress item to the list.
          setStatus("loading");
          setLoadingMessage(e.data.data);
          break;

        case "initiate":
          setProgressItems((prev) => [...prev, e.data]);
          break;

        case "progress":
          // Model file progress: update one of the progress items.
          setProgressItems((prev) =>
            prev.map((item) => {
              if (item.file === e.data.file) {
                return { ...item, ...e.data };
              }
              return item;
            }),
          );
          break;

        case "done":
          // Model file loaded: remove the progress item from the list.
          setProgressItems((prev) =>
            prev.filter((item) => item.file !== e.data.file),
          );
          break;

        case "ready":
          // Pipeline ready: the worker is ready to accept messages.
          setStatus("ready");
          break;

        case "start":
          {
            // Start generation
            setMessages((prev) => [
              ...prev,
              { role: "assistant", content: "" },
            ]);
          }
          break;

        case "update":
          {
            // Generation update: update the output text.
            // Parse messages
            const { output, tps, numTokens, state } = e.data;
            setTps(tps);
            setNumTokens(numTokens);
            setMessages((prev) => {
              const cloned = [...prev];
              const last = cloned.at(-1);
              const data = {
                ...last,
                content: last.content + output,
              };
              if (data.answerIndex === undefined && state === "answering") {
                // When state changes to answering, we set the answerIndex
                data.answerIndex = last.content.length;
              }
              cloned[cloned.length - 1] = data;
              return cloned;
            });
          }
          break;

        case "complete":
          // Generation complete: re-enable the "Generate" button
          setIsRunning(false);
          break;

        case "error":
          setError(e.data.data);
          break;
      }
    };

    const onErrorReceived = (e) => {
      console.error("Worker error:", e);
    };

    // Attach the callback function as an event listener.
    worker.current.addEventListener("message", onMessageReceived);
    worker.current.addEventListener("error", onErrorReceived);

    // Define a cleanup function for when the component is unmounted.
    return () => {
      worker.current.removeEventListener("message", onMessageReceived);
      worker.current.removeEventListener("error", onErrorReceived);
    };
  }, []);

  // Send the messages to the worker thread whenever the `messages` state changes.
  useEffect(() => {
    if (messages.filter((x) => x.role === "user").length === 0) {
      // No user messages yet: do nothing.
      return;
    }
    if (messages.at(-1).role === "assistant") {
      // Do not update if the last message is from the assistant
      return;
    }
    setTps(null);
    worker.current.postMessage({ type: "generate", data: messages });
  }, [messages, isRunning]);

  useEffect(() => {
    if (!chatContainerRef.current || !isRunning) return;
    const element = chatContainerRef.current;
    if (
      element.scrollHeight - element.scrollTop - element.clientHeight <
      STICKY_SCROLL_THRESHOLD
    ) {
      element.scrollTop = element.scrollHeight;
    }
  }, [messages, isRunning]);

  return IS_WEBGPU_AVAILABLE ? (
    <div className="flex flex-col h-screen mx-auto items justify-end text-slate-800 dark:text-slate-200 bg-gradient-to-b from-slate-50 via-white to-blue-50 dark:from-gray-900 dark:via-gray-800 dark:to-blue-900/20 transition-colors duration-300">
      {status === null && messages.length === 0 && (
        <div className="h-full overflow-auto scrollbar-thin flex justify-center items-center flex-col relative bg-white/50 dark:bg-gray-900/50 backdrop-blur-sm">
          <div className="flex flex-col items-center mb-1 max-w-[400px] text-center animate-fade-in">
            <div className="text-7xl mb-4 flex items-center gap-3 animate-bounce-slow">
              <span className="animate-pulse">🤔</span>
              <span className="animate-pulse animation-delay-200">💭</span>
              <span className="animate-pulse animation-delay-400">🧠</span>
            </div>
            <h1 className={`text-5xl font-bold mb-3 ${GRADIENT_CLASSES} bg-clip-text text-transparent transition-all duration-500`}>
              DeepSeek-R1 WebGPU
            </h1>
            <div className="text-8xl mb-4 group perspective">
              <div className="relative transform-gpu transition-all duration-500 animate-float group-hover:scale-110">
                <span className="absolute inset-0 animate-shimmer-fast bg-gradient-to-r from-transparent via-white/30 to-transparent rounded-full blur-sm"></span>
                <span className="relative inline-block animate-shimmer bg-gradient-to-r from-blue-400 via-cyan-300 to-blue-400 bg-clip-text text-transparent filter drop-shadow-glow hover:drop-shadow-glow-strong">
                  🐳
                </span>
                <div className="absolute inset-0 bg-gradient-to-r from-blue-500/20 via-cyan-400/20 to-blue-500/20 rounded-full blur-md animate-pulse-slow"></div>
              </div>
            </div>
            <h2 className="font-semibold text-slate-700 dark:text-slate-300">
              A next-generation reasoning model that runs locally in your browser with WebGPU acceleration.
            </h2>
            <p className="text-sm mt-2 text-slate-500 dark:text-slate-400">
              Forked from{" "}
              <a
                href="https://github.com/huggingface/transformers.js-examples/tree/main/deepseek-r1-webgpu"
                target="_blank"
                rel="noreferrer"
                className="font-medium text-blue-500 dark:text-blue-400 hover:text-blue-600 dark:hover:text-blue-300 transition-colors"
              >
                Xenova & Deepseek-R1
              </a>
              {" "}- Frontend by{" "}
              <a
                href="https://github.com/Reidmen"
                target="_blank"
                rel="noreferrer"
                className="font-medium text-blue-500 dark:text-blue-400 hover:text-blue-600 dark:hover:text-blue-300 transition-colors"
              >
                Reidmen
              </a>
            </p>
          </div>

          <div className="flex flex-col items-center px-4 animate-fade-in-up">
            <p className="max-w-[510px] mb-4 text-slate-600 dark:text-slate-400">
              <br />
              You are about to load{" "}
              <a
                href="https://huggingface.co/onnx-community/DeepSeek-R1-Distill-Qwen-1.5B-ONNX"
                target="_blank"
                rel="noreferrer"
                className="font-medium underline"
              >
                DeepSeek-R1-Distill-Qwen-1.5B
              </a>
              , a 1.5B parameter reasoning LLM. 
              Everything runs entirely in your browser with{" "}
       
            </p>

            {error && (
              <div className="text-red-500 text-center mb-4 p-4 bg-red-50 dark:bg-red-900/20 rounded-xl animate-fade-in">
                <p className="font-semibold mb-2">
                  Unable to load model due to the following error:
                </p>
                <p className="text-sm mb-3 font-mono bg-red-100/50 dark:bg-red-900/40 p-2 rounded">{error}</p>
                
                <div className="text-sm text-red-600 dark:text-red-300">
                  {error.includes("WebGPU") ? (
                    <p>
                      💡 This app requires WebGPU support. Please:
                      <ul className="list-disc list-inside mt-1">
                        <li>Enable WebGPU flags in your browser settings</li>
                      </ul>
                    </p>
                  ) : error.includes("memory") ? (
                    <p>
                      💡 Memory issue detected. Try:
                      <ul className="list-disc list-inside mt-1">
                        <li>Closing other browser tabs</li>
                        <li>Refreshing the page</li>
                        <li>Using a device with more available memory</li>
                      </ul>
                    </p>
                  ) : error.includes("fetch") ? (
                    <p>
                      💡 Network error detected. Please:
                      <ul className="list-disc list-inside mt-1">
                        <li>Check your internet connection</li>
                        <li>Verify if Hugging Face is accessible</li>
                        <li>Try refreshing the page</li>
                      </ul>
                    </p>
                  ) : (
                    <p>
                      💡 Try refreshing the page or check the browser console for more details.
                      If the issue persists, please report it on our{" "}
                      <a 
                        href="https://github.com/huggingface/transformers.js-examples/issues"
                        target="_blank"
                        rel="noreferrer"
                        className="underline hover:text-red-400"
                      >
                        GitHub issues
                      </a>.
                    </p>
                  )}
                </div>
              </div>
            )}

            <button
              className={`border px-8 py-3 rounded-2xl ${GRADIENT_CLASSES} text-white font-medium shadow-lg shadow-blue-500/30 dark:shadow-blue-800/30 hover:shadow-xl hover:shadow-blue-500/40 dark:hover:shadow-blue-800/40 disabled:from-slate-400 disabled:to-slate-500 disabled:shadow-none transition-all duration-300 transform hover:scale-105 disabled:hover:scale-100`}
              onClick={() => {
                worker.current.postMessage({ type: "load" });
                setStatus("loading");
              }}
              disabled={status !== null || error !== null}
            >
              Load model
            </button>
          </div>
        </div>
      )}
      {status === "loading" && (
        <div className="w-full max-w-[500px] text-left mx-auto p-6 bottom-0 mt-auto bg-white/90 dark:bg-gray-800/90 backdrop-blur-md rounded-2xl shadow-xl animate-fade-in">
          <p className="text-center mb-1 text-slate-700 dark:text-slate-300">{loadingMessage}</p>
          {progressItems.map(({ file, progress, total }, i) => (
            <Progress
              key={i}
              text={file}
              percentage={progress}
              total={total}
              className="bg-gradient-to-r from-blue-500 to-indigo-500"
            />
          ))}
        </div>
      )}

      {status === "ready" && (
        <div
          ref={chatContainerRef}
          className="overflow-y-auto scrollbar-thin w-full flex flex-col items-center h-full bg-gradient-to-b from-transparent to-blue-50/30 dark:to-blue-900/20"
        >
          <Chat messages={messages} />
          {messages.length === 0 && (
            <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-4 p-6 animate-fade-in-up">
              {EXAMPLES.map((msg, i) => (
                <div
                  key={i}
                  className={`border border-gray-200 dark:border-gray-700 rounded-xl p-4 bg-white/80 dark:bg-gray-800/80 backdrop-blur-sm cursor-pointer ${CARD_HOVER_CLASSES} transition-all duration-300`}
                  onClick={() => onEnter(msg)}
                >
                  {msg}
                </div>
              ))}
            </div>
          )}
          <p className="text-center text-sm min-h-6 text-gray-500 dark:text-gray-300">
            {tps && messages.length > 0 && (
              <>
                {!isRunning && (
                  <span>
                    Generated {numTokens} tokens in{" "}
                    {(numTokens / tps).toFixed(2)} seconds&nbsp;&#40;
                  </span>
                )}
                {
                  <>
                    <span className="font-medium text-center mr-1 text-black dark:text-white">
                      {tps.toFixed(2)}
                    </span>
                    <span className="text-gray-500 dark:text-gray-300">
                      tokens/second
                    </span>
                  </>
                }
                {!isRunning && (
                  <>
                    <span className="mr-1">&#41;.</span>
                    <span
                      className="underline cursor-pointer"
                      onClick={() => {
                        worker.current.postMessage({ type: "reset" });
                        setMessages([]);
                      }}
                    >
                      Reset
                    </span>
                  </>
                )}
              </>
            )}
          </p>
        </div>
      )}

      <div className="mt-4 border-0 rounded-2xl w-[700px] max-w-[90%] mx-auto relative mb-4 bg-white/90 dark:bg-gray-800/90 backdrop-blur-md shadow-lg hover:shadow-xl transition-all duration-300">
        <div className="flex items-center w-full p-2">
          <textarea
            ref={textareaRef}
            className="w-full px-4 py-3 rounded-xl bg-transparent border-none outline-none text-slate-800 dark:text-slate-200 placeholder-slate-400 dark:placeholder-slate-500 disabled:text-slate-400 disabled:placeholder-slate-300 resize-none disabled:cursor-not-allowed focus:ring-2 focus:ring-blue-500/50 transition-all duration-300"
            placeholder="Type your message..."
            type="text"
            rows={1}
            value={input}
            disabled={status !== "ready"}
            title={status === "ready" ? "Model is ready" : "Model not loaded yet"}
            onKeyDown={(e) => {
              if (input.length > 0 && !isRunning && e.key === "Enter" && !e.shiftKey) {
                e.preventDefault();
                onEnter(input);
              }
            }}
            onInput={(e) => setInput(e.target.value)}
          />
          <div className="flex-shrink-0 px-2">
            {isRunning ? (
              <button 
                onClick={onInterrupt}
                className="p-2 rounded-lg hover:bg-red-50 dark:hover:bg-red-900/30 text-red-500 dark:text-red-400 transition-all duration-300 hover:scale-110"
                title="Stop generating"
              >
                <StopIcon className="h-5 w-5" />
              </button>
            ) : (
              <button
                onClick={() => input.length > 0 && onEnter(input)}
                disabled={input.length === 0}
                className={`p-2 rounded-lg transition-all duration-300 ${
                  input.length > 0
                    ? `${GRADIENT_CLASSES} text-white shadow-md shadow-blue-500/30 dark:shadow-blue-800/30 hover:shadow-lg hover:scale-110`
                    : "bg-gray-100 dark:bg-gray-700 text-gray-400 dark:text-gray-600 cursor-not-allowed"
                }`}
                title={input.length > 0 ? "Send message" : "Please type a message"}
              >
                <ArrowRightIcon className="h-5 w-5" />
              </button>
            )}
          </div>
        </div>
      </div>

      <p className="text-xs text-slate-400 dark:text-slate-500 text-center mb-3 animate-pulse">
        Disclaimer: Generated content may be inaccurate or false.
      </p>
    </div>
  ) : (
    <div className="fixed w-screen h-screen bg-black z-10 bg-opacity-[92%] text-white text-2xl font-semibold flex justify-center items-center text-center">
      WebGPU is not supported
      <br />
      by this browser :&#40;
    </div>
  );
}

export default App;