frontend/src/components/SimpleChat.jsx

import { useState, useEffect, useRef, useMemo } 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, onSend, isLoading }) => {
  const [input, setInput] = useState('');
  const containerRef = useRef(null);
  const anchorRef = useRef(null); // <- will be a tiny zero-height anchor BEFORE the bubble
  const textareaRef = useRef(null);

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

  const handleKeyDown = (e) => {
    if (e.key === 'Enter' && !e.shiftKey) {
      e.preventDefault();
      handleSubmit(e);
    }
  };

  // Auto-resize textarea
  useEffect(() => {
    if (textareaRef.current) {
      textareaRef.current.style.height = 'auto';
      textareaRef.current.style.height = Math.min(textareaRef.current.scrollHeight, 120) + 'px';
    }
  }, [input]);

  // Since messages are now filtered to current chunk only, use the last message for anchoring
  const { anchorIndex, firstInChunkIndex } = useMemo(() => {
    const lastIndex = messages.length > 0 ? messages.length - 1 : -1;
    const firstIndex = messages.length > 0 ? 0 : -1;
    return { anchorIndex: lastIndex, firstInChunkIndex: firstIndex };
  }, [messages]);

  // Scroll by scrolling the ZERO-HEIGHT anchor into view AFTER layout commits.
  const scrollAfterLayout = () => {
    requestAnimationFrame(() => {
      requestAnimationFrame(() => {
        if (anchorRef.current && containerRef.current) {
          // Calculate position relative to container instead of using scrollIntoView
          const anchorTop = anchorRef.current.offsetTop;
          containerRef.current.scrollTo({ top: anchorTop, behavior: 'smooth' });
        } else if (containerRef.current) {
          // fallback: go to top
          containerRef.current.scrollTo({ top: 0, behavior: 'smooth' });
        }
      });
    });
  };

  // When chunk changes, try to pin.
  useEffect(() => {
    if (anchorIndex !== -1) {
      scrollAfterLayout();
    } else if (containerRef.current) {
      requestAnimationFrame(() => containerRef.current.scrollTo({ top: 0, behavior: 'smooth' }));
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [currentChunkIndex, anchorIndex]);

  // New messages: pin the new anchor after layout
  useEffect(() => {
    if (anchorIndex !== -1) scrollAfterLayout();
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [messages.length, anchorIndex]);

  return (
    <div className="flex flex-col h-full min-h-0">
      <div
        ref={containerRef}
        className="flex-1 min-h-0 overflow-y-auto p-4 flex flex-col space-y-3"
      >
        {messages.map((message, idx) => {
          const isAnchor = idx === anchorIndex;

          // Render a zero-height anchor just BEFORE the bubble for the anchor index.
          if (isAnchor) {
            return (
              <div key={idx} className="flex flex-col">
                {/* <-- ZERO-HEIGHT anchor: deterministic top-of-message alignment */}
                <div ref={anchorRef} style={{ height: 0, margin: 0, padding: 0 }} />

                <div className={`flex ${message.role === 'user' ? 'justify-end' : 'justify-start'}`}>
                  <div
                    className={`max-w-[90%] p-3 rounded-lg ${
                      message.role === 'user'
                        ? 'bg-gray-100 text-white'
                        : 'bg-white text-gray-900'
                    }`}
                  >
                    <ReactMarkdown
                      remarkPlugins={[remarkMath]}
                      rehypePlugins={[rehypeRaw, rehypeKatex]}
                      components={getChatMarkdownComponents()}
                    >
                      {message.content}
                    </ReactMarkdown>
                  </div>
                </div>

                {isLoading && (
                  <div className="flex justify-start mt-3">
                    <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>
                )}

                {/* filler to push remaining whitespace below the pinned message */}
                <div className="flex-1" />
              </div>
            );
          }

          // Non-anchor message: render normally
          return (
            <div
              key={idx}
              className={`flex ${message.role === 'user' ? 'justify-end' : 'justify-start'}`}
            >
              <div
                className={`max-w-[90%] p-3 rounded-lg ${
                            message.role === 'user'
                              ? 'bg-gray-100 text-white'
                              : 'bg-white text-gray-900'
                          }`}
              >
                <ReactMarkdown
                  remarkPlugins={[remarkMath]}
                  rehypePlugins={[rehypeRaw, rehypeKatex]}
                  components={getChatMarkdownComponents()}
                >
                  {message.content}
                </ReactMarkdown>
              </div>
            </div>
          );
        })}

        {/* if no messages in chunk yet, render typing+filler */}
        {firstInChunkIndex === -1 && (
          <div className="flex flex-col">
            {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 className="flex-1" />
          </div>
        )}
      </div>

      {/* Input with auto-resize textarea */}
      <form onSubmit={handleSubmit} className="p-4 border-t">
        <div className="flex space-x-2 items-end">
          <textarea
            ref={textareaRef}
            value={input}
            onChange={(e) => setInput(e.target.value)}
            onKeyDown={handleKeyDown}
            placeholder="Type your message... (Enter to send, Shift+Enter for new line)"
            disabled={isLoading}
            rows={1}
            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 resize-none overflow-hidden"
            style={{ minHeight: '42px' }}
          />
          <button
            type="submit"
            disabled={!input.trim() || isLoading}
            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;