frontend/src/components/DocumentViewer.jsx

import { useState, useRef, useEffect } from 'react';
import { PdfLoader, PdfHighlighter, useHighlightContainerContext, TextHighlight, AreaHighlight } from 'react-pdf-highlighter-extended';
import * as pdfjs from "pdfjs-dist";

// Tell pdf.js to use the local worker file
pdfjs.GlobalWorkerOptions.workerSrc = '/pdf.worker.min.js';

// Copy exact example from documentation
const MyHighlightContainer = () => {
  const {
    highlight, // The highlight being rendered
    viewportToScaled, // Convert a highlight position to platform agnostic coords (useful for saving edits)
    screenshot, // Screenshot a bounding rectangle
    isScrolledTo, // Whether the highlight has been auto-scrolled to
    highlightBindings, // Whether the highlight has been auto-scrolled to
  } = useHighlightContainerContext();

  const isTextHighlight = !Boolean(
    highlight.content && highlight.content.image
  );

  const component = isTextHighlight ? (
    <TextHighlight
      isScrolledTo={isScrolledTo}
      highlight={highlight}
    />
  ) : (
    <AreaHighlight
      isScrolledTo={isScrolledTo}
      highlight={highlight}
      onChange={(boundingRect) => {
        const edit = {
          position: {
            boundingRect: viewportToScaled(boundingRect),
            rects: [],
          },
          content: {
            image: screenshot(boundingRect),
          },
        };
      }}
      bounds={highlightBindings.textLayer}
    />
  );

  return component;
};

const DocumentViewer = ({ selectedFile, documentData, onPageChange, preloadedHighlights = null, currentChunkIndex = null, onDocumentReady = null, isChunkLoading = null }) => {
    const [highlights, setHighlights] = useState([]);
    const [pdfUrl, setPdfUrl] = useState(null);
    const [zoom, setZoom] = useState(1);
    
    /** Refs for PdfHighlighter utilities */
    const highlighterUtilsRef = useRef();
    const documentReadyCalledRef = useRef(false);
    
    // Function to scroll to a specific chunk's highlight
    const scrollToChunk = (chunkIndex) => {
        if (highlighterUtilsRef.current && preloadedHighlights) {
            const chunkHighlights = preloadedHighlights[chunkIndex];
            if (chunkHighlights && chunkHighlights.length > 0) {
                const firstHighlightInChunk = chunkHighlights[0];
                highlighterUtilsRef.current.scrollToHighlight(firstHighlightInChunk);
            }
        }
    };

    // Function to scroll to the first highlight (for backwards compatibility)
    const scrollToFirstChunk = () => {
        scrollToChunk(0);
    };

    const zoomWithCenter = (zoomDelta) => {
        const container = document.querySelector('.PdfHighlighter');
        if (!container) {
            setZoom(prev => prev + zoomDelta);
            return;
        }

        const scrollLeft = container.scrollLeft;
        const scrollTop = container.scrollTop;
        const containerWidth = container.clientWidth;
        const containerHeight = container.clientHeight;
        
        // Calculate center point before zoom
        const centerX = scrollLeft + containerWidth / 2;
        const centerY = scrollTop + containerHeight / 2;
        
        setZoom(prev => {
            const newZoom = prev + zoomDelta;
            const zoomRatio = newZoom / prev;
            
            // Use requestAnimationFrame for smoother transition
            requestAnimationFrame(() => {
                const newScrollLeft = centerX * zoomRatio - containerWidth / 2;
                const newScrollTop = centerY * zoomRatio - containerHeight / 2;
                container.scrollTo({
                    left: newScrollLeft,
                    top: newScrollTop,
                    behavior: 'auto'
                });
            });
            
            return newZoom;
        });
    };

    const zoomIn = () => {
        if (zoom < 4) {
            zoomWithCenter(0.1);
        }
    };

    const zoomOut = () => {
        if (zoom > 0.2) {
            zoomWithCenter(-0.1);
        }
    };

    const resetZoom = () => {
        setZoom(1);
    };

    // Call onDocumentReady only once when utils become available
    useEffect(() => {
        if (onDocumentReady && !documentReadyCalledRef.current && highlighterUtilsRef.current) {
            documentReadyCalledRef.current = true;
            onDocumentReady({ scrollToFirstChunk });
        }
    }, [onDocumentReady, scrollToFirstChunk, highlighterUtilsRef.current]);

    // Utility function to normalize highlight data
    const normalizeHighlight = (highlightData) => {
        // Ensure the highlight has the required structure
        if (!highlightData.id || !highlightData.position || !highlightData.content) {
            return null;
        }
        
        return {
            id: highlightData.id,
            position: highlightData.position,
            content: highlightData.content
        };
    };

    // Convert File object to URL for PdfLoader
    useEffect(() => {
        if (selectedFile) {
            if (typeof selectedFile === 'string') {
                setPdfUrl(selectedFile);
            } else if (selectedFile instanceof File) {
                const url = URL.createObjectURL(selectedFile);
                setPdfUrl(url);
                return () => URL.revokeObjectURL(url);
            }
        } else {
            setPdfUrl(null);
        }
    }, [selectedFile]);

    // Load preloaded highlights when component mounts or when currentChunkIndex changes
    useEffect(() => {
        if (preloadedHighlights) {
            let highlightsToLoad = [];
            
            if (currentChunkIndex !== null && currentChunkIndex !== undefined && preloadedHighlights[currentChunkIndex]) {
                // Load highlights for specific chunk
                highlightsToLoad = preloadedHighlights[currentChunkIndex];
            } else if (Array.isArray(preloadedHighlights)) {
                // Load all highlights if it's an array
                highlightsToLoad = preloadedHighlights;
            } else if (typeof preloadedHighlights === 'object') {
                // If it's an object without chunkIndex, take all values
                highlightsToLoad = Object.values(preloadedHighlights).flat();
            }
            
            // Normalize and filter valid highlights
            const validHighlights = highlightsToLoad
                .map(normalizeHighlight)
                .filter(Boolean);
            
            console.log(`🎨 Loading ${validHighlights.length} preloaded highlights${currentChunkIndex !== null ? ` for chunk ${currentChunkIndex}` : ''}`);
            setHighlights(validHighlights);
        } else {
            // Clear highlights if no preloaded data
            setHighlights([]);
        }
    }, [preloadedHighlights, currentChunkIndex]);

    // Auto-scroll to current chunk when currentChunkIndex changes (only on navigation, not during streaming)
    useEffect(() => {
        // Only auto-scroll if we have highlighter utils and this is a valid chunk navigation
        // Don't auto-scroll during streaming (when isChunkLoading is true for currentChunkIndex)
        if (highlighterUtilsRef.current && 
            currentChunkIndex !== null && 
            currentChunkIndex !== undefined && 
            currentChunkIndex >= 0 &&
            (!isChunkLoading || !isChunkLoading(currentChunkIndex))) {
            // Small delay to ensure highlights are loaded
            setTimeout(() => {
                scrollToChunk(currentChunkIndex);
            }, 200);
        }
    }, [currentChunkIndex, isChunkLoading]); // Only depend on currentChunkIndex, not preloadedHighlights

    // Handle selection - log coordinates and add debugging
    const handleSelection = (selection) => {
        console.log("🎯 SELECTION MADE! Full selection object:", selection);
        console.log("📍 Position:", selection.position);
        console.log("📝 Content:", selection.content);
        console.log("🔍 Type:", selection.type);
        
        const newHighlight = {
            id: `highlight_${Date.now()}`,
            position: selection.position,
            content: selection.content
        };
        
        console.log("✅ Adding highlight:", newHighlight);
        setHighlights(prev => [...prev, newHighlight]);
    };

    // Additional debugging handlers
    const handleCreateGhost = (ghost) => {
        console.log("👻 Ghost highlight created:", ghost);
    };

    const handleRemoveGhost = (ghost) => {
        console.log("❌ Ghost highlight removed:", ghost);
    };

    if (!selectedFile || !pdfUrl) {
        return (
            <div className="bg-white rounded-lg shadow-sm flex items-center justify-center h-full">
                <div className="text-center text-gray-500">
                    <p>No PDF selected</p>
                </div>
            </div>
        );
    }
    return (
        <div className="bg-white rounded-lg shadow-sm flex flex-col relative" style={{ width: '100%', height: '100%' }}>
            <div className="flex justify-between items-center p-2 border-b">
                <h2>{documentData?.filename || 'Document'}</h2>
                <div className="flex items-center gap-2">
                    <button 
                        title="Zoom out" 
                        onClick={zoomOut}
                        className="px-2 py-1 border rounded hover:bg-gray-100"
                        disabled={zoom <= 0.2}
                    >
                        -
                    </button>
                    <button 
                        title="Reset zoom" 
                        onClick={resetZoom}
                        className="px-2 py-1 border rounded hover:bg-gray-100 text-sm"
                        disabled={zoom === 1}
                    >
                        {(zoom * 100).toFixed(0)}%
                    </button>
                    <button 
                        title="Zoom in" 
                        onClick={zoomIn}
                        className="px-2 py-1 border rounded hover:bg-gray-100"
                        disabled={zoom >= 4}
                    >
                        +
                    </button>
                </div>
            </div>
            
            <div style={{ height: '500px' }}>
                <PdfLoader document={pdfUrl} workerSrc='/pdf.worker.min.js'>
                    {(pdfDocument) => (
                        <PdfHighlighter
                            enableAreaSelection={(event) => event.altKey}
                            pdfDocument={pdfDocument}
                            pdfScaleValue={zoom}
                            utilsRef={(_pdfHighlighterUtils) => {
                                highlighterUtilsRef.current = _pdfHighlighterUtils;
                            }}
                            highlights={highlights}
                            onSelection={handleSelection}
                            onCreateGhostHighlight={handleCreateGhost}
                            onRemoveGhostHighlight={handleRemoveGhost}
                        >
                            <MyHighlightContainer />
                        </PdfHighlighter>
                    )}
                </PdfLoader>
            </div>
        </div>
    );
};

export default DocumentViewer;