import React from 'react';
import { observer } from 'mobx-react';
import { AssistantMessage, Citation, MessageContent, TEXT_TYPE } from './types';
import Markdown from 'react-markdown';

interface MessageComponentProps {
    message: AssistantMessage;
    index: number;
    onFollowUpClick: (question: string) => void;
    onCitationClick: (citation: Citation) => void;
    updateMessageCitations: (index: number, citations: Citation[]) => void;
}

const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message, index, onFollowUpClick, onCitationClick, updateMessageCitations }) {
    const renderContent = (item: MessageContent) => {
        const i = item.index;
        if (item.type === TEXT_TYPE.GROUNDED) {
            const citation_ids = item.citation_ids || [];
            return (
                <span key={i} className="grounded-text">
                    {item.text}
                    {citation_ids.map((id, idx) => {
                        const citation = message.citations?.find(c => c.citation_id === id);
                        if (!citation) return null;
                        return (
                            <button
                                key={idx}
                                className="citation-button"
                                onClick={() => onCitationClick(citation)}
                                style={{
                                    display: 'inline-flex',
                                    alignItems: 'center',
                                    justifyContent: 'center',
                                    width: '20px',
                                    height: '20px',
                                    borderRadius: '50%',
                                    border: 'none',
                                    background: '#ff6347',
                                    color: 'white',
                                    fontSize: '12px',
                                    fontWeight: 'bold',
                                    cursor: 'pointer',
                                    margin: '0 2px',
                                    padding: 0,
                                }}>
                                {idx + 1}
                            </button>
                        );
                    })}
                </span>
            );
        } else {
            return <span key={i}>{item.text}</span>;
        }
    };

    return (
        <div className={`message ${message.role}`}>
            <div>{message.content && message.content.map(messageFragment => <React.Fragment key={messageFragment.index}>{renderContent(messageFragment)}</React.Fragment>)}</div>
            {message.follow_up_questions && message.follow_up_questions.length > 0 && (
                <div className="follow-up-questions">
                    <h4>Follow-up Questions:</h4>
                    {message.follow_up_questions.map((question, idx) => (
                        <button key={idx} className="follow-up-button" onClick={() => onFollowUpClick(question)}>
                            {question}
                        </button>
                    ))}
                </div>
            )}
        </div>
    );
};

export default observer(MessageComponentBox);