frontend/src/components/ChunkPanel.jsx

import ReactMarkdown from 'react-markdown';
import remarkMath from 'remark-math';
import rehypeKatex from 'rehype-katex';
import rehypeRaw from 'rehype-raw';
import { getChatMarkdownComponents, getTitleMarkdownComponents } from '../utils/markdownComponents.jsx';
import SimpleChat from './SimpleChat.jsx';
import ChunkLoadingTips from './ChunkLoadingTips.jsx';
import React, { useEffect } from 'react';
import { createTextStreamResponse } from 'ai';

const ChunkPanel = ({
    documentData,
    currentChunkIndex,
    showChat,
    updateGlobalChatHistory,
    getGlobalChatHistory,
    addMessageToChunk,
    getCurrentChunkMessages,
    hasChunkMessages,
    markChunkUnderstood,
    skipChunk,
    goToPrevChunk,
    streamResponse,
    isChunkLoading
}) => {

    const chatMarkdownComponents = getChatMarkdownComponents();
    const titleMarkdownComponents = getTitleMarkdownComponents();

    // Generate greeting for chunks that don't have messages yet
    // Only for initial chunk (0) and when not transitioning
    useEffect(() => {
        if (documentData && showChat && !hasChunkMessages(currentChunkIndex) && currentChunkIndex === 0) {
            console.log("🤖 Triggering greeting generation for chunk 0");
            generateGreetingStreaming();
        }
    }, [documentData?.chunks?.length, showChat, currentChunkIndex]); // More stable dependencies

    const updateLastAssistantMessage = (delta) => {
        const allMessages = getGlobalChatHistory();

        const currentChunkMessages = allMessages.filter(msg => msg.chunkIndex === currentChunkIndex);
        const lastAssistantInChunk = [...currentChunkMessages].reverse().find(msg => msg.role === 'assistant');

        if (!lastAssistantInChunk) {
            console.warn("No assistant message found for current chunk — adding new one.");
            addMessageToChunk({ role: 'assistant', content: delta }, currentChunkIndex);
            return;
        }

        const updatedMessages = allMessages.map(msg => {
            if (msg === lastAssistantInChunk) {
                return { ...msg, content: msg.content + (typeof delta === "string" ? delta : delta?.content || "") };
            }
            return msg;
        });

        updateGlobalChatHistory(updatedMessages);
    };

    const generateGreetingStreaming = async () => {
        const requestBody = JSON.stringify({
            messages: [],
            currentChunk: documentData?.chunks?.[currentChunkIndex]?.text || '',
            document: documentData ? JSON.stringify(documentData) : '',
        });

        streamResponse(requestBody, false);
    };


    const handleSendStreaming = async (text) => {
        const userMessage = { role: 'user', content: text, chunkIndex: currentChunkIndex };
        addMessageToChunk(userMessage, currentChunkIndex);

        // Build the messages array manually to include the user message immediately
        const messagesWithUserMessage = [...getGlobalChatHistory(), userMessage];

        const requestBody = JSON.stringify({
            messages: messagesWithUserMessage,
            currentChunk: documentData?.chunks?.[currentChunkIndex]?.text || '',
            document: documentData ? JSON.stringify(documentData) : ''
        });

        streamResponse(requestBody, false);
    };


    return (
        <>
            {/* Chunk Header */}
            <div className="px-6 py-4 flex-shrink-0 bg-white rounded-t-lg border-b border-gray-200 z-10 flex items-center justify-between">
                        <div className="flex items-center flex-1">
                            {/* Previous Chunk Button */}
                            <button
                                onClick={goToPrevChunk}
                                disabled={currentChunkIndex === 0}
                                className="mr-3 p-2 rounded-full bg-gray-100 hover:bg-gray-200 text-gray-600 transition-colors duration-200 disabled:opacity-50 disabled:cursor-not-allowed"
                                title="Go to previous chunk"
                            >
                                <svg 
                                    className="w-5 h-5" 
                                    fill="currentColor" 
                                    viewBox="0 0 20 20"
                                >
                                    <path 
                                        fillRule="evenodd" 
                                        d="M12.707 5.293a1 1 0 010 1.414L9.414 10l3.293 3.293a1 1 0 01-1.414 1.414l-4-4a1 1 0 010-1.414l4-4a1 1 0 011.414 0z" 
                                        clipRule="evenodd" 
                                    />
                                </svg>
                            </button>
                            
                            {/* Chunk Title */}
                            <div className="font-semibold text-xl text-gray-900 text-left">
                                <ReactMarkdown
                                    remarkPlugins={[remarkMath]}
                                    rehypePlugins={[rehypeRaw, rehypeKatex]}
                                    components={titleMarkdownComponents}
                                >
                                    {documentData?.chunks?.[currentChunkIndex]?.topic || "Loading..."}
                                </ReactMarkdown>
                            </div>
                        </div>
                        
                        <div className="flex items-center gap-2">
                            {/* Skip Button */}
                            <button
                                onClick={skipChunk}
                                className="p-2 rounded-full bg-gray-100 hover:bg-gray-200 text-gray-600 transition-colors duration-200"
                                title="Skip this chunk"
                            >
                                <svg 
                                    className="w-5 h-5" 
                                    fill="currentColor" 
                                    viewBox="0 0 20 20"
                                >
                                    <path 
                                        fillRule="evenodd" 
                                        d="M7.293 14.707a1 1 0 010-1.414L10.586 10 7.293 6.707a1 1 0 011.414-1.414l4 4a1 1 0 010 1.414l-4 4a1 1 0 01-1.414 0z" 
                                        clipRule="evenodd" 
                                    />
                                    <path 
                                        fillRule="evenodd" 
                                        d="M12.293 14.707a1 1 0 010-1.414L15.586 10l-3.293-3.293a1 1 0 011.414-1.414l4 4a1 1 0 010 1.414l-4 4a1 1 0 01-1.414 0z" 
                                        clipRule="evenodd" 
                                    />
                                </svg>
                            </button>
                            
                            {/* Understood Button */}
                            <button
                                onClick={markChunkUnderstood}
                                className="p-2 rounded-full bg-green-100 hover:bg-green-200 text-green-600 transition-colors duration-200"
                                title="Mark chunk as understood"
                            >
                                <svg 
                                    className="w-5 h-5" 
                                    fill="currentColor" 
                                    viewBox="0 0 20 20"
                                >
                                    <path 
                                        fillRule="evenodd" 
                                        d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z" 
                                        clipRule="evenodd" 
                                    />
                                </svg>
                            </button>
                        </div>

            </div>

            {/* Chat Interface - Always shown when showChat is true */}
            {showChat && (
                <div className="relative flex-1 overflow-hidden">
                    <SimpleChat
                        messages={getCurrentChunkMessages()}
                        currentChunkIndex={currentChunkIndex}
                        onSend={handleSendStreaming}
                        isLoading={isChunkLoading(currentChunkIndex)}
                    />
                </div>
            )}
        </>
    );
};

export default ChunkPanel;