client/src/components/MessageBubble.tsx

import { cn } from "@/lib/utils";
import { User, Bot, Copy, Check, RotateCcw, Pencil, Minimize2 } from "lucide-react";
import { useState } from "react";
import { ToolCallCard } from "./ToolCallCard";
import { ThinkingBlock, extractThinkingBlocks } from "./ThinkingBlock";
import { Streamdown } from "streamdown";
import type { ChatMessage, ContentSegment } from "@/hooks/useChat";

export function MessageBubble({
  message,
  onRetry,
  onEdit,
}: {
  message: ChatMessage;
  onRetry?: () => void;
  onEdit?: (content: string) => void;
}) {
  const [copied, setCopied] = useState(false);
  const [isEditing, setIsEditing] = useState(false);
  const [editValue, setEditValue] = useState(message.content);
  const isUser = message.role === "user";
  const isSystem = message.role === "system";

  const copyContent = () => {
    if (message.content) {
      navigator.clipboard.writeText(message.content);
      setCopied(true);
      setTimeout(() => setCopied(false), 2000);
    }
  };

  const handleEdit = () => {
    setIsEditing(true);
    setEditValue(message.content);
  };

  const submitEdit = () => {
    if (onEdit && editValue.trim()) {
      onEdit(editValue.trim());
    }
    setIsEditing(false);
  };

  if (isSystem) {
    // Detect compact report (matches original format_compact_report)
    const isCompactReport = message.content.startsWith("Compact\n");
    if (isCompactReport) {
      return (
        <div className="flex justify-center my-3">
          <div className="flex items-center gap-2 text-xs text-amber-400/80 bg-amber-500/10 border border-amber-500/20 px-4 py-2 rounded-lg max-w-xl font-mono">
            <Minimize2 className="size-3.5 shrink-0" />
            <pre className="whitespace-pre text-[11px] leading-relaxed">{message.content}</pre>
          </div>
        </div>
      );
    }
    return (
      <div className="flex justify-center my-2">
        <div className="text-xs text-muted-foreground bg-accent/30 px-3 py-1.5 rounded-full max-w-xl">
          <Streamdown>{message.content}</Streamdown>
        </div>
      </div>
    );
  }

  // ─── Segmented rendering for assistant messages ───────────────────
  // If we have contentSegments, render them in order (interleaved text + tools).
  // Otherwise fall back to legacy rendering (all tools then all text).
  const hasSegments = !isUser && message.contentSegments && message.contentSegments.length > 0;
  const toolCalls = message.toolCalls || [];

  const renderSegmentedContent = () => {
    if (!message.contentSegments) return null;

    return message.contentSegments.map((segment: ContentSegment, i: number) => {
      if (segment.type === "text") {
        const text = segment.content?.trim();
        if (!text) return null;
        return (
          <div
            key={`seg-text-${i}`}
            className="text-sm leading-relaxed text-foreground/90 prose prose-invert prose-sm max-w-none mt-1.5"
          >
            <Streamdown>{segment.content}</Streamdown>
          </div>
        );
      }
      if (segment.type === "tool") {
        const tool = toolCalls[segment.toolIndex];
        if (!tool) return null;
        return (
          <div key={`seg-tool-${i}`} className="mt-1.5">
            <ToolCallCard tool={tool} />
          </div>
        );
      }
      return null;
    });
  };

  const renderLegacyContent = () => {
    return (
      <>
        {/* Tool calls — show BEFORE message text */}
        {toolCalls.length > 0 && (
          <div className="mt-2 space-y-1">
            {toolCalls.map((tool) => (
              <ToolCallCard key={tool.id} tool={tool} />
            ))}
          </div>
        )}

        {/* Message text — after tool calls */}
        {isEditing ? (
          <div className="inline-block text-left max-w-[85%]">
            <textarea
              value={editValue}
              onChange={(e) => setEditValue(e.target.value)}
              className="w-full bg-input border border-border rounded-lg px-3 py-2 text-sm font-mono outline-none focus:border-primary resize-none min-h-[60px]"
              autoFocus
              onKeyDown={(e) => {
                if (e.key === "Enter" && !e.shiftKey) {
                  e.preventDefault();
                  submitEdit();
                }
                if (e.key === "Escape") setIsEditing(false);
              }}
            />
            <div className="flex gap-1.5 mt-1 justify-end">
              <button
                onClick={() => setIsEditing(false)}
                className="text-xs px-2 py-1 rounded bg-muted text-muted-foreground hover:bg-accent"
              >
                Cancel
              </button>
              <button
                onClick={submitEdit}
                className="text-xs px-2 py-1 rounded bg-primary text-primary-foreground hover:bg-primary/90"
              >
                Save
              </button>
            </div>
          </div>
        ) : message.content ? (
          <div
            className={cn(
              "text-sm leading-relaxed",
              isUser
                ? "bg-primary/10 rounded-2xl rounded-tr-sm px-4 py-2.5 inline-block text-left max-w-[85%]"
                : "text-foreground/90 prose prose-invert prose-sm max-w-none",
              toolCalls.length > 0 ? "mt-2" : ""
            )}
          >
            {isUser ? (
              <span className="whitespace-pre-wrap">{message.content}</span>
            ) : (
              <Streamdown>{message.content}</Streamdown>
            )}
            {message.isStreaming && <span className="streaming-cursor" />}
          </div>
        ) : null}
      </>
    );
  };

  return (
    <div
      className={cn(
        "group flex gap-3 py-4 px-4",
        isUser ? "flex-row-reverse" : ""
      )}
    >
      {/* Avatar */}
      <div
        className={cn(
          "size-7 rounded-lg shrink-0 flex items-center justify-center mt-0.5",
          isUser
            ? "bg-primary/20 text-primary"
            : "bg-accent text-accent-foreground"
        )}
      >
        {isUser ? <User className="size-4" /> : <Bot className="size-4" />}
      </div>

      {/* Content */}
      <div className={cn("flex-1 min-w-0", isUser ? "text-right" : "")}>
        <div
          className={cn(
            "flex items-center gap-2 mb-1",
            isUser ? "justify-end" : ""
          )}
        >
          <span className="text-xs font-medium text-muted-foreground">
            {isUser ? "You" : "Claw"}
          </span>
          {message.model && (
            <span className="text-[10px] text-muted-foreground/60 font-mono">
              {message.model}
            </span>
          )}
          {/* Action buttons */}
          <div className="opacity-0 group-hover:opacity-100 transition-opacity flex items-center gap-1">
            {message.content && (
              <button
                onClick={copyContent}
                className="text-muted-foreground hover:text-foreground transition-colors"
                title="Copy"
              >
                {copied ? (
                  <Check className="size-3" />
                ) : (
                  <Copy className="size-3" />
                )}
              </button>
            )}
            {isUser && onEdit && (
              <button
                onClick={handleEdit}
                className="text-muted-foreground hover:text-foreground transition-colors"
                title="Edit"
              >
                <Pencil className="size-3" />
              </button>
            )}
            {!isUser && onRetry && (
              <button
                onClick={onRetry}
                className="text-muted-foreground hover:text-foreground transition-colors"
                title="Retry"
              >
                <RotateCcw className="size-3" />
              </button>
            )}
          </div>
        </div>

        {/* Thinking blocks — show first (model reasoning) */}
        {!isUser && message.thinkingBlocks && message.thinkingBlocks.length > 0 && (
          <div className="mt-1">
            {message.thinkingBlocks.map((block: any, i: number) => (
              <ThinkingBlock key={i} thinking={block.thinking} durationMs={block.durationMs} />
            ))}
          </div>
        )}

        {/* ─── Main content area ─── */}
        {hasSegments ? (
          // New: Interleaved rendering — text and tools in the order they arrived
          <div className="mt-1">
            {renderSegmentedContent()}
            {message.isStreaming && <span className="streaming-cursor" />}
          </div>
        ) : (
          // Legacy: all tools first, then all text
          renderLegacyContent()
        )}

        {/* Editing UI for user messages */}
        {isUser && isEditing && (
          <div className="inline-block text-left max-w-[85%]">
            <textarea
              value={editValue}
              onChange={(e) => setEditValue(e.target.value)}
              className="w-full bg-input border border-border rounded-lg px-3 py-2 text-sm font-mono outline-none focus:border-primary resize-none min-h-[60px]"
              autoFocus
              onKeyDown={(e) => {
                if (e.key === "Enter" && !e.shiftKey) {
                  e.preventDefault();
                  submitEdit();
                }
                if (e.key === "Escape") setIsEditing(false);
              }}
            />
            <div className="flex gap-1.5 mt-1 justify-end">
              <button
                onClick={() => setIsEditing(false)}
                className="text-xs px-2 py-1 rounded bg-muted text-muted-foreground hover:bg-accent"
              >
                Cancel
              </button>
              <button
                onClick={submitEdit}
                className="text-xs px-2 py-1 rounded bg-primary text-primary-foreground hover:bg-primary/90"
              >
                Save
              </button>
            </div>
          </div>
        )}

        {/* Usage info */}
        {message.cost !== undefined && message.cost > 0 && (
          <div className="mt-1 flex items-center gap-2 text-[10px] text-muted-foreground/60">
            {message.promptTokens && (
              <span>{message.promptTokens.toLocaleString()} in</span>
            )}
            {message.completionTokens && (
              <span>{message.completionTokens.toLocaleString()} out</span>
            )}
            <span>${message.cost.toFixed(6)}</span>
          </div>
        )}
      </div>
    </div>
  );
}