import { useState, useEffect, useRef } from 'react';
import ReactMarkdown from 'react-markdown';
import remarkMath from 'remark-math';
import rehypeKatex from 'rehype-katex';
import rehypeRaw from 'rehype-raw';
import { getChatMarkdownComponents } from '../utils/markdownComponents.jsx';


const SimpleChat = ({ messages, currentChunkIndex, canEdit, onSend, isLoading }) => {
  const [input, setInput] = useState('');
  const messagesEndRef = useRef(null);
  const currentChunkStartRef = useRef(null);

  const handleSubmit = (e) => {
    e.preventDefault();
    if (!input.trim() || isLoading || !canEdit) return;
    onSend(input.trim());
    setInput('');
  };

  // Scroll to current chunk's first message when chunk changes
  useEffect(() => {
    if (currentChunkStartRef.current) {
      currentChunkStartRef.current.scrollIntoView({ behavior: 'smooth', block: 'start' });
    }
  }, [currentChunkIndex]);

  // Find the first message of the current chunk
  const currentChunkFirstMessageIndex = messages.findIndex(msg => msg.chunkIndex === currentChunkIndex);

  return (
    <div className="flex flex-col h-full">
      {/* Messages */}
      <div className="flex-1 overflow-y-auto p-4 space-y-3">
        {messages.map((message, idx) => {
          const isCurrentChunk = message.chunkIndex === currentChunkIndex;
          const isFirstOfCurrentChunk = idx === currentChunkFirstMessageIndex;
          
          return (
            <div
              key={idx}
              ref={isFirstOfCurrentChunk ? currentChunkStartRef : null}
              className={`flex ${message.role === 'user' ? 'justify-end' : 'justify-start'}`}
            >
              <div
                className={`max-w-[90%] p-3 rounded-lg transition-opacity ${
                  message.role === 'user'
                    ? `bg-gray-100 text-white ${isCurrentChunk ? 'opacity-100' : 'opacity-40'}`
                    : `bg-white text-gray-900 ${isCurrentChunk ? 'opacity-100' : 'opacity-40'}`
                }`}
              >
                <ReactMarkdown
                    remarkPlugins={[remarkMath]}
                    rehypePlugins={[rehypeRaw, rehypeKatex]}
                    components={getChatMarkdownComponents()}
                  >
                    {message.content}
                  </ReactMarkdown>
              </div>
            </div>
          );
        })}

        {isLoading && (
          <div className="flex justify-start">
            <div className="bg-gray-100 p-3 rounded-lg">
              <div className="flex space-x-1">
                <div className="w-2 h-2 bg-gray-400 rounded-full animate-bounce"></div>
                <div
                  className="w-2 h-2 bg-gray-400 rounded-full animate-bounce"
                  style={{ animationDelay: '0.1s' }}
                ></div>
                <div
                  className="w-2 h-2 bg-gray-400 rounded-full animate-bounce"
                  style={{ animationDelay: '0.2s' }}
                ></div>
              </div>
            </div>
          </div>
        )}
      </div>

      {/* Input */}
      <form onSubmit={handleSubmit} className="p-4 border-t">
        <div className="flex space-x-2">
          <input
            type="text"
            value={input}
            onChange={(e) => setInput(e.target.value)}
            placeholder={canEdit ? "Type your message..." : "This chunk is completed - navigation only"}
            disabled={isLoading || !canEdit}
            className="flex-1 px-3 py-2 border rounded-lg focus:outline-none focus:ring-2 focus:ring-blue-500 disabled:bg-gray-100 disabled:text-gray-500"
          />
          <button
            type="submit"
            disabled={!input.trim() || isLoading || !canEdit}
            className="px-4 py-2 bg-blue-500 text-white rounded-lg hover:bg-blue-600 disabled:bg-gray-300 disabled:cursor-not-allowed"
          >
            {isLoading ? '...' : 'Send'}
          </button>
        </div>
      </form>
    </div>
  );
};

export default SimpleChat;