src/chat/MessageContent.tsx

import { Button } from "@theme";
import cn from "@utils/classnames.ts";
import { Lightbulb } from "lucide-react";
import { useEffect, useState } from "react";
import showdown from "showdown";

const converter = new showdown.Converter();
interface ParsedContent {
  thinkContent: string | null;
  afterContent: string;
  isThinking: boolean;
}

function parseThinkTags(content: string): ParsedContent {
  const openTagIndex = content.indexOf("<think>");

  if (openTagIndex === -1) {
    return {
      thinkContent: null,
      afterContent: content,
      isThinking: false,
    };
  }

  const closeTagIndex = content.indexOf("</think>");

  if (closeTagIndex === -1) {
    return {
      thinkContent: content.slice(openTagIndex + 7),
      afterContent: "",
      isThinking: true,
    };
  }

  return {
    thinkContent: content.slice(openTagIndex + 7, closeTagIndex),
    afterContent: content.slice(closeTagIndex + 8),
    isThinking: false,
  };
}

const PROSE_CLASS_NAME =
  "prose dark:prose-invert prose-headings:font-semibold prose-headings:mt-4 prose-headings:mb-2 prose-h3:text-base prose-p:my-2 prose-ul:my-2 prose-li:my-0";

export default function MessageContent({ content }: { content: string }) {
  const [showThinking, setShowThinking] = useState(false);
  const [thinkingTime, setThinkingTime] = useState(0);
  const [forceThinkingComplete, setForceThinkingComplete] = useState(false);
  const parsed = parseThinkTags(content);

  // Detect when content stops changing and force thinking to complete
  useEffect(() => {
    if (parsed.isThinking) {
      const timeout = setTimeout(() => {
        setForceThinkingComplete(true);
      }, 1000);

      return () => clearTimeout(timeout);
    } else {
      setForceThinkingComplete(false);
    }
  }, [content, parsed.isThinking]);

  useEffect(() => {
    if (parsed.isThinking && !forceThinkingComplete) {
      const startTime = Date.now();
      const interval = setInterval(() => {
        setThinkingTime((Date.now() - startTime) / 1000);
      }, 100);
      return () => clearInterval(interval);
    }
  }, [parsed.isThinking, forceThinkingComplete]);

  const isThinking = parsed.isThinking && !forceThinkingComplete;

  if (!parsed.thinkContent) {
    return (
      <div
        className={cn(PROSE_CLASS_NAME, "max-w-none")}
        dangerouslySetInnerHTML={{
          __html: converter.makeHtml(content),
        }}
      />
    );
  }

  return (
    <div className="space-y-2">
      <div>
        <Button
          variant="ghost"
          color="mono"
          size="xs"
          onClick={() => setShowThinking(!showThinking)}
          className="-ml-2 opacity-50"
          loading={isThinking}
          iconLeft={<Lightbulb />}
          notDisabledWhileLoading
        >
          {isThinking ? "Thinking for" : "Thought for"}{" "}
          {thinkingTime.toFixed(1)}s ({showThinking ? "Hide" : "Show"})
        </Button>
        {showThinking && (
          <p className="my-4 max-w-none border-l-1 border-gray-300 pl-4 text-sm font-medium text-gray-600 dark:border-gray-700 dark:text-gray-400">
            {parsed.thinkContent}
          </p>
        )}
      </div>
      {parsed.afterContent && (
        <div
          className={cn(PROSE_CLASS_NAME, "max-w-none")}
          dangerouslySetInnerHTML={{
            __html: converter.makeHtml(parsed.afterContent),
          }}
        />
      )}
    </div>
  );
}