import React from 'react';
import ChatError from './Messages/ErrorMessage';
import useRenderTracker from '../../hooks/useRenderTracker';
import { useAssistantThinkingSummary, type Message } from '../../contexts/ChatContext';
import UserMessage from './Messages/UserMessage';
import StaticAssistantMessage from './Messages/StaticAssistantMessage';
import StreamingAssistantMessage from './Messages/StreamingAssistantMessage';

interface MessageWrapperProps {
  message: Message;
  turnKey?: string;
  onRetry: (id: string) => void;
  onDismiss: (id: string) => void;
}

const MessageWrapper: React.FC<MessageWrapperProps> = ({ message, turnKey, onRetry, onDismiss }) => {
  const thinkingSummary = useAssistantThinkingSummary(message.role === 'assistant' ? turnKey : undefined);
  
  useRenderTracker('Message', { 
    message,
    thinkingSummary,
    onRetry,
    onDismiss,
  });

  // Don't render empty assistant messages
  if (message.role === 'assistant' && !message.content.trim()) {
    return null;

  } else if (message.role === 'error') {
    return (
      <ChatError
        errorMessage={message.content}
        onRetry={() => onRetry(message.userMessageId || message.message_id)}
        onDismiss={() => onDismiss(message.userMessageId || message.message_id)}
      />
    );

  } else if (message.role === 'user') {
    return (
      <div className={`chat-message user`}>
        <UserMessage
          content={message.content}
        />
      </div>
    );
  } else if (message.role === 'assistant') {
    return (
      <div className={`chat-message assistant`}>
        <StaticAssistantMessage
          content={message.content}
          thinkingSummary={thinkingSummary || undefined}
        />
      </div>
    );
  } else if (message.role === 'assistant-streaming') {
    return (
      <div className={`chat-message assistant`}>
        <StreamingAssistantMessage/>
      </div>
    );
  }
};

// Add custom comparison function to prevent unnecessary re-renders
export default React.memo(MessageWrapper, (prevProps, nextProps) => {
  return prevProps.message.message_id === nextProps.message.message_id &&
         prevProps.message.content === nextProps.message.content;
});