src/components/MessageBubble.tsx

import { useState, useRef, useEffect, useCallback } from "react";
import { Streamdown } from "streamdown";
import { createMathPlugin } from "@streamdown/math";
import { Pencil, X, Check, RotateCcw, Copy, ClipboardCheck, Gauge } from "lucide-react";

import { useLLM } from "../hooks/useLLM";
import type { ChatMessage } from "../hooks/LLMContext";

const math = createMathPlugin({ singleDollarTextMath: true });

interface MessageBubbleProps {
  msg: ChatMessage;
  index: number;
  isStreaming?: boolean;
  isGenerating: boolean;
}

const MATH_COMMANDS: { prefix: string; args: number }[] = [
  { prefix: "\\boxed{", args: 1 },
  { prefix: "\\text{", args: 1 },
  { prefix: "\\textbf{", args: 1 },
  { prefix: "\\mathbf{", args: 1 },
  { prefix: "\\mathrm{", args: 1 },
  { prefix: "\\frac{", args: 2 },
];

function skipBraceGroup(content: string, start: number): number {
  let depth = 1;
  let j = start;
  while (j < content.length && depth > 0) {
    if (content[j] === "{") depth++;
    else if (content[j] === "}") depth--;
    j++;
  }
  return j;
}

function wrapLatexMath(content: string): string {
  let result = "";
  let i = 0;
  let mathContext: null | "$" | "$$" = null;

  while (i < content.length) {
    const cmd = !mathContext ? MATH_COMMANDS.find((candidate) => content.startsWith(candidate.prefix, i)) : undefined;

    if (cmd) {
      let j = skipBraceGroup(content, i + cmd.prefix.length);
      for (let argIndex = 1; argIndex < cmd.args; argIndex++) {
        if (content[j] === "{") {
          j = skipBraceGroup(content, j + 1);
        }
      }

      const expr = content.slice(i, j);
      result += "$" + expr + "$";
      i = j;
    } else if (content[i] === "$") {
      const isDouble = content[i + 1] === "$";
      const token = isDouble ? "$$" : "$";

      if (mathContext === token) mathContext = null;
      else if (!mathContext) mathContext = token;

      result += token;
      i += token.length;
    } else {
      result += content[i];
      i++;
    }
  }

  return result;
}

function prepareForMathDisplay(content: string): string {
  return wrapLatexMath(
    content
      .replace(/(?<!\\)\\\[/g, "$$$$")
      .replace(/\\\]/g, "$$$$")
      .replace(/(?<!\\)\\\(/g, "$$$$")
      .replace(/\\\)/g, "$$$$"),
  );
}

export function MessageBubble({ msg, index, isStreaming, isGenerating }: MessageBubbleProps) {
  const { editMessage, retryMessage } = useLLM();
  const isUser = msg.role === "user";
  const isAssistant = msg.role === "assistant";

  const [isEditing, setIsEditing] = useState(false);
  const [editValue, setEditValue] = useState(msg.content);
  const [copied, setCopied] = useState(false);
  const textareaRef = useRef<HTMLTextAreaElement>(null);

  const handleCopy = useCallback(async () => {
    await navigator.clipboard.writeText(msg.content);
    setCopied(true);
    setTimeout(() => setCopied(false), 2000);
  }, [msg.content]);

  const handleRetry = useCallback(() => {
    retryMessage(index);
  }, [retryMessage, index]);

  useEffect(() => {
    if (isEditing && textareaRef.current) {
      textareaRef.current.focus();
      textareaRef.current.style.height = "auto";
      textareaRef.current.style.height = textareaRef.current.scrollHeight + "px";
    }
  }, [isEditing]);

  const handleEdit = useCallback(() => {
    setEditValue(msg.content);
    setIsEditing(true);
  }, [msg.content]);

  const handleCancel = useCallback(() => {
    setIsEditing(false);
    setEditValue(msg.content);
  }, [msg.content]);

  const handleSave = useCallback(() => {
    const trimmed = editValue.trim();
    if (!trimmed) return;
    setIsEditing(false);
    editMessage(index, trimmed);
  }, [editValue, editMessage, index]);

  const handleKeyDown = useCallback(
    (event: React.KeyboardEvent) => {
      if (event.key === "Escape") handleCancel();
      if (event.key === "Enter" && !event.shiftKey) {
        event.preventDefault();
        handleSave();
      }
    },
    [handleCancel, handleSave],
  );

  const displayTps = !isStreaming && isAssistant && msg.tps ? msg.tps : null;

  if (isEditing) {
    return (
      <div className="flex justify-end">
        <div className="w-full max-w-[80%] flex flex-col gap-2">
          <textarea
            ref={textareaRef}
            value={editValue}
            onChange={(event) => {
              setEditValue(event.target.value);
              event.target.style.height = "auto";
              event.target.style.height = event.target.scrollHeight + "px";
            }}
            onKeyDown={handleKeyDown}
            className="w-full py-3 px-4 border border-line rounded-2xl bg-[rgba(255,255,255,0.04)] text-text font-body text-[0.93rem] leading-[1.5] resize-none focus:outline-[1px] focus:outline-[rgba(157,224,255,0.44)] focus:border-[rgba(157,224,255,0.44)]"
            rows={1}
          />
          <div className="flex justify-end gap-2">
            <button
              onClick={handleCancel}
              className="flex items-center gap-1.5 py-1.5 px-3 border border-line rounded-[10px] bg-transparent text-text-muted text-[0.8rem] font-medium transition-[color,background-color] duration-[180ms] ease-[ease] hover:text-text hover:bg-[rgba(255,255,255,0.06)]"
            >
              <X size={12} />
              Cancel
            </button>
            <button
              onClick={handleSave}
              disabled={!editValue.trim()}
              className="flex items-center gap-1.5 py-1.5 px-3 border-0 rounded-[10px] bg-[rgba(157,224,255,0.18)] text-accent-strong text-[0.8rem] font-medium transition-[background-color,opacity] duration-[180ms] ease-[ease] hover:bg-[rgba(157,224,255,0.28)] disabled:opacity-40"
            >
              <Check size={12} />
              Update
            </button>
          </div>
        </div>
      </div>
    );
  }

  return (
    <>
      <div className={`group flex items-start gap-2 ${isUser ? "justify-end" : "justify-start"}`}>
        {isUser && !isGenerating && (
          <button
            onClick={handleEdit}
            className="mt-3 p-1 rounded-md text-text-muted bg-transparent opacity-0 group-hover:opacity-100 transition-[opacity,color] duration-[180ms] ease-[ease] hover:text-text"
            title="Edit message"
          >
            <Pencil size={14} />
          </button>
        )}

        <div
          className={`max-w-[80%] py-3.5 px-[18px] text-[0.93rem] leading-[1.55] border border-line bg-panel backdrop-blur-[24px] ${
            isUser
              ? "rounded-[24px_24px_6px_24px] border-[rgba(157,224,255,0.25)] bg-[rgba(157,224,255,0.08)] whitespace-pre-wrap"
              : "rounded-[24px_24px_24px_6px]"
          }`}
        >
          {msg.content ? (
            isUser ? (
              msg.content
            ) : (
              <Streamdown plugins={{ math }} parseIncompleteMarkdown={false} isAnimating={isStreaming}>
                {prepareForMathDisplay(msg.content)}
              </Streamdown>
            )
          ) : !isUser && !isStreaming ? (
            <p className="italic text-text-muted">No response</p>
          ) : null}
        </div>

        {!isUser && !isStreaming && !isGenerating && (
          <div className="flex items-center gap-1 mt-3 opacity-0 group-hover:opacity-100 transition-opacity duration-[180ms] ease-[ease]">
            {msg.content && (
              <button
                onClick={handleCopy}
                className="p-1 rounded-md text-text-muted bg-transparent transition-[color,background-color] duration-[180ms] ease-[ease] hover:text-text hover:bg-[rgba(255,255,255,0.08)]"
                title="Copy response"
              >
                {copied ? <ClipboardCheck size={14} /> : <Copy size={14} />}
              </button>
            )}
            <button
              onClick={handleRetry}
              className="p-1 rounded-md text-text-muted bg-transparent transition-[color,background-color] duration-[180ms] ease-[ease] hover:text-text hover:bg-[rgba(255,255,255,0.08)]"
              title="Retry"
            >
              <RotateCcw size={14} />
            </button>
          </div>
        )}
      </div>

      {displayTps && (
        <div className="flex items-center gap-1 mt-1 ml-3 text-text-muted font-mono text-[0.7rem] tabular-nums">
          <Gauge size={11} />
          <span>{displayTps.toFixed(1)} tok/s</span>
        </div>
      )}
    </>
  );
}