import React from 'react';
import { observer } from 'mobx-react';
import { AssistantMessage, Citation } 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 = (content: string) => {
        if (!message.citations || message.citations.length === 0) {
            return content;
        }

        const parts = [];
        let lastIndex = 0;

        message.citations.forEach((citation, idx) => {
            const location = citation.text_location;
            const textBefore = content.slice(lastIndex, location);
            const citationButton = (
                <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>
            );
            parts.push(textBefore, citationButton);
            lastIndex = location;
        });

        parts.push(content.slice(lastIndex));

        return parts;
    };

    return (
        <div className={`message ${message.role}`}>
            <div>{renderContent(message.text_content)}</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);
