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,
} from "lucide-react";

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

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

interface MessageBubbleProps {
  msg: ChatMessage;
  index: number;
  isStreaming?: boolean;
  thinkingSeconds?: number;
  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,
  thinkingSeconds,
  isGenerating,
}: MessageBubbleProps) {
  const { editMessage, retryMessage } = useLLM();
  const isUser = msg.role === "user";
  const isThinking = !!isStreaming && !msg.content;

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

  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 rounded-xl border border-[#2a2a2a] bg-[#111111] px-4 py-3 text-sm text-[#f0f0f0] placeholder-[#888888] focus:border-[#c9a84c] focus:outline-none focus:ring-1 focus:ring-[#c9a84c] resize-none shadow-sm"
            rows={1}
          />
          <div className="flex justify-end gap-2">
            <button
              onClick={handleCancel}
              className="flex items-center gap-1.5 rounded-lg px-3 py-1.5 text-xs font-medium text-[#888888] hover:text-[#f0f0f0] border border-[#2a2a2a] hover:bg-[#1a1a1a] transition-colors cursor-pointer"
            >
              <X className="h-3 w-3" />
              Cancel
            </button>
            <button
              onClick={handleSave}
              disabled={!editValue.trim()}
              className="flex items-center gap-1.5 rounded-lg bg-[#76b900] px-3 py-1.5 text-xs font-medium text-black hover:bg-[#8fd400] disabled:opacity-40 transition-colors cursor-pointer"
            >
              <Check className="h-3 w-3" />
              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 opacity-0 group-hover:opacity-100 transition-opacity text-[#888888] hover:text-[#f0f0f0] cursor-pointer"
          title="Edit message"
        >
          <Pencil className="h-3.5 w-3.5" />
        </button>
      )}

      <div
        className={`max-w-[80%] rounded-2xl px-4 py-3 text-sm leading-relaxed ${
          isUser
            ? "bg-[#111111] text-[#f0f0f0] rounded-br-md whitespace-pre-wrap border border-[#c9a84c33]"
            : "bg-[#111111] text-[#f0f0f0] rounded-bl-md border border-[#2a2a2a] shadow-sm"
        }`}
      >
        {!isUser && msg.reasoning && (
          <ReasoningBlock
            reasoning={msg.reasoning}
            isThinking={isThinking}
            thinkingSeconds={msg.thinkingSeconds ?? thinkingSeconds ?? 0}
          />
        )}

        {msg.content ? (
          isUser ? (
            msg.content
          ) : (
            <Streamdown
              plugins={{ math }}
              parseIncompleteMarkdown={false}
              isAnimating={isStreaming}
            >
              {prepareForMathDisplay(msg.content)}
            </Streamdown>
          )
        ) : !isUser && !isStreaming ? (
          <p className="italic text-[#888888]">No response</p>
        ) : null}
      </div>

      {!isUser && !isStreaming && !isGenerating && (
        <div className="mt-3 flex items-center gap-1 opacity-0 group-hover:opacity-100 transition-opacity">
          {msg.content && (
            <button
              onClick={handleCopy}
              className="rounded-md p-1 text-[#888888] hover:text-[#f0f0f0] hover:bg-[#1a1a1a] transition-colors cursor-pointer"
              title="Copy response"
            >
              {copied ? (
                <ClipboardCheck className="h-3.5 w-3.5" />
              ) : (
                <Copy className="h-3.5 w-3.5" />
              )}
            </button>
          )}
          <button
            onClick={handleRetry}
            className="rounded-md p-1 text-[#6d6d6d] hover:text-black hover:bg-[#f5f5f5] transition-colors cursor-pointer"
            title="Retry"
          >
            <RotateCcw className="h-3.5 w-3.5" />
          </button>
        </div>
      )}
    </div>
  );
}