Update app.py

app.py

@@ -442,7 +442,7 @@ class DownloadManager:
             return None
 
     async def force_download_viewonly(self, file_info, save_path):
-        """Last-resort method to download view-only Google Drive files"""
+        """Last-resort method to download view-only Google Drive files - improved for multi-page PDFs"""
         try:
             # Extract file ID from URL
             file_id = None
@@ -469,11 +469,12 @@ class DownloadManager:
             # Launch a new browser context with higher resolution
             browser = await self.playwright.chromium.launch(
                 headless=True,
-                args=['--no-sandbox', '--disable-setuid-sandbox', '--disable-dev-shm-usage']
+                args=['--no-sandbox', '--disable-setuid-sandbox', '--disable-dev-shm-usage', '--disable-web-security']
             )
             
+            # Use a larger viewport for better quality
             context = await browser.new_context(
-                viewport={'width': 1600, 'height': 1200},
+                viewport={'width': 1920, 'height': 1080},
                 user_agent=get_random_user_agent(),
                 device_scale_factor=2.0  # Higher resolution for better quality
             )
@@ -482,135 +483,310 @@ class DownloadManager:
             
             # Navigate to the file
             try:
+                logger.info(f"Opening view-only file: https://drive.google.com/file/d/{file_id}/view")
                 await page.goto(f"https://drive.google.com/file/d/{file_id}/view", 
                                wait_until='networkidle', 
-                               timeout=60000)
+                               timeout=90000)  # Longer timeout for large PDFs
                 
                 # Wait for content to load fully
                 await page.wait_for_timeout(5000)
                 
-                # Check if it's a PDF
+                # Detect if it's a PDF
                 is_pdf = await page.query_selector('embed[type="application/pdf"]') is not None
                 
                 if is_pdf:
-                    # For PDFs: Screenshot each page approach
-                    logger.info("Detected PDF, using page-by-page screenshot approach")
+                    # For PDFs: Multi-page capture approach
+                    logger.info("Detected PDF, using multi-page capture approach")
                     
-                    # Scroll through document to ensure all pages are loaded
+                    # First, try to find the viewer container
+                    viewer_container = await page.query_selector('.drive-viewer-paginated-scrollable')
+                    
+                    if not viewer_container:
+                        logger.warning("Could not find standard PDF viewer container, trying alternatives")
+                        viewer_container = await page.query_selector('.drive-viewer-content') or \
+                                          await page.query_selector('#drive-pdf-viewer') or \
+                                          await page.query_selector('.drive-viewer')
+                    
+                    if not viewer_container:
+                        # Take a single screenshot as fallback
+                        logger.warning("Could not find any PDF viewer container, using fallback")
+                        screenshot_path = os.path.join(tempfile.gettempdir(), "gdrive_pdf_fallback.png")
+                        await page.screenshot(path=screenshot_path, full_page=True)
+                        
+                        # Convert to PDF
+                        from PIL import Image
+                        from reportlab.pdfgen import canvas as pdf_canvas
+                        img = Image.open(screenshot_path)
+                        width, height = img.size
+                        c = pdf_canvas.Canvas(save_path, pagesize=(width, height))
+                        c.drawImage(screenshot_path, 0, 0, width, height)
+                        c.save()
+                        os.remove(screenshot_path)
+                        return save_path
+                    
+                    # Scroll through to load all pages first
+                    logger.info("Pre-loading all PDF pages...")
                     await page.evaluate("""
-                        async function scrollDocument() {
+                        async function preloadAllPages() {
                             const delay = ms => new Promise(resolve => setTimeout(resolve, ms));
                             const container = document.querySelector('.drive-viewer-paginated-scrollable');
                             if (!container) return;
                             
-                            // First scroll to bottom to load all pages
+                            // Scroll to bottom to force all pages to load
+                            const initialScroll = container.scrollTop;
                             container.scrollTo(0, container.scrollHeight);
-                            await delay(2000);
+                            await delay(3000); // Wait for loading
                             
-                            // Then back to top
+                            // Scroll back to top
                             container.scrollTo(0, 0);
                             await delay(1000);
                         }
-                        return scrollDocument();
+                        return preloadAllPages();
                     """)
                     
-                    # Count pages
+                    # Count visible pages - critical step that needs to be fixed
                     page_count = await page.evaluate("""
                         () => {
+                            // Try multiple selectors for pages
                             const pages = document.querySelectorAll('.drive-viewer-paginated-page');
-                            return pages.length;
+                            if (pages.length > 0) return pages.length;
+                            
+                            // Alternative selectors if standard one fails
+                            const altPages = document.querySelectorAll('.drive-viewer-page');
+                            if (altPages.length > 0) return altPages.length;
+                            
+                            // Try to find page numbers in navigation
+                            const pageNav = document.querySelector('.drive-viewer-paginated-counter');
+                            if (pageNav) {
+                                const text = pageNav.textContent || '';
+                                const match = text.match(/(\d+)\s*\/\s*(\d+)/);
+                                if (match && match[2]) return parseInt(match[2]);
+                            }
+                            
+                            return 0; // Fallback
                         }
                     """)
                     
+                    # If no pages found but we know it's a PDF, manually check for page counter
                     if page_count == 0:
-                        logger.warning("No pages found, trying alternative method")
-                        # Take a screenshot of the entire page
-                        temp_dir = tempfile.mkdtemp()
-                        screenshot_path = os.path.join(temp_dir, "page.png")
-                        await page.screenshot(path=screenshot_path, full_page=True)
-                        
-                        # Convert screenshot to PDF
-                        from PIL import Image
-                        from reportlab.pdfgen import canvas as pdf_canvas
-                        img = Image.open(screenshot_path)
-                        width, height = img.size
-                        c = pdf_canvas.Canvas(save_path, pagesize=(width, height))
-                        c.drawImage(screenshot_path, 0, 0, width, height)
-                        c.save()
-                        
-                        # Clean up
-                        os.remove(screenshot_path)
-                        os.rmdir(temp_dir)
-                    else:
-                        # Create temp directory for page screenshots
-                        temp_dir = tempfile.mkdtemp()
-                        screenshots = []
+                        # Try to find the page counter text and extract total pages
+                        page_counter_text = await page.evaluate("""
+                            () => {
+                                const elements = Array.from(document.querySelectorAll('*'));
+                                for (const el of elements) {
+                                    const text = el.textContent || '';
+                                    if (text.match(/\d+\s*\/\s*\d+/)) return text;
+                                }
+                                return '';
+                            }
+                        """)
                         
-                        # Take screenshot of each page
-                        for i in range(page_count):
-                            # Scroll to page
-                            await page.evaluate(f"""
-                                async () => {{
-                                    const pages = document.querySelectorAll('.drive-viewer-paginated-page');
-                                    if (pages.length <= {i}) return;
-                                    const delay = ms => new Promise(resolve => setTimeout(resolve, ms));
-                                    pages[{i}].scrollIntoView();
+                        if page_counter_text:
+                            match = re.search(r'(\d+)\s*\/\s*(\d+)', page_counter_text)
+                            if match and match.group(2):
+                                page_count = int(match.group(2))
+                                logger.info(f"Detected {page_count} pages from page counter")
+                    
+                    # If we still have no page count, default to a reasonable number
+                    if page_count == 0:
+                        logger.warning("Could not detect page count, defaulting to 50 pages to be safe")
+                        page_count = 50  # Try to capture up to 50 pages by default
+                    
+                    logger.info(f"Found {page_count} pages in PDF")
+                    
+                    # Create a temporary directory for screenshots
+                    temp_dir = tempfile.mkdtemp()
+                    screenshots = []
+                    
+                    # Function to scroll to a specific page and take a screenshot
+                    async def capture_page(page_num):
+                        # Scroll to the page
+                        success = await page.evaluate(f"""
+                            async function scrollToPage(pageNum) {{
+                                const delay = ms => new Promise(resolve => setTimeout(resolve, ms));
+                                
+                                // Try multiple selectors for pages
+                                const pages = document.querySelectorAll('.drive-viewer-paginated-page');
+                                if (pages.length > 0 && pageNum < pages.length) {{
+                                    pages[pageNum].scrollIntoView({{behavior: 'instant', block: 'center'}});
                                     await delay(500);
+                                    return true;
                                 }}
-                            """)
-                            
-                            # Take screenshot
-                            screenshot_path = os.path.join(temp_dir, f"page_{i+1}.png")
-                            
-                            # Position page for best screenshot
-                            await page.evaluate(f"""
-                                () => {{
-                                    const pages = document.querySelectorAll('.drive-viewer-paginated-page');
-                                    const page = pages[{i}];
-                                    const viewer = document.querySelector('.drive-viewer-paginated-scrollable');
-                                    if (page && viewer) {{
-                                        // Center the page in the viewport
-                                        const rect = page.getBoundingClientRect();
-                                        viewer.scrollBy(0, rect.top - 100);
+                                
+                                // Alternative: try to use page navigation buttons
+                                const pageInput = document.querySelector('input[aria-label="Page"]');
+                                if (pageInput) {{
+                                    // Set page number in input
+                                    const nativeInputValueSetter = Object.getOwnPropertyDescriptor(window.HTMLInputElement.prototype, "value").set;
+                                    nativeInputValueSetter.call(pageInput, {page_num + 1});
+                                    
+                                    // Dispatch events
+                                    const ev1 = new Event('input', {{ bubbles: true }});
+                                    const ev2 = new Event('change', {{ bubbles: true }});
+                                    pageInput.dispatchEvent(ev1);
+                                    pageInput.dispatchEvent(ev2);
+                                    
+                                    // Press Enter to navigate
+                                    const keyEvent = new KeyboardEvent('keydown', {{
+                                        key: 'Enter',
+                                        code: 'Enter',
+                                        keyCode: 13,
+                                        which: 13,
+                                        bubbles: true
+                                    }});
+                                    pageInput.dispatchEvent(keyEvent);
+                                    
+                                    await delay(1000); // Wait for navigation
+                                    return true;
+                                }}
+                                
+                                // Alternative: use page selector dropdown if available
+                                const pageSelector = document.querySelector('.drive-viewer-paginated-page-selector');
+                                if (pageSelector) {{
+                                    pageSelector.click();
+                                    await delay(300);
+                                    
+                                    // Find and click the specific page option
+                                    const options = document.querySelectorAll('.drive-viewer-paginated-page-selector-option');
+                                    if (options.length > pageNum) {{
+                                        options[pageNum].click();
+                                        await delay(1000);
+                                        return true;
                                     }}
                                 }}
-                            """)
+                                
+                                return false;
+                            }}
+                            return scrollToPage({page_num});
+                        """)
+                        
+                        if not success:
+                            # Alternative: Try using the page navigation buttons
+                            logger.info(f"Using alternative navigation for page {page_num + 1}")
                             
-                            await page.screenshot(path=screenshot_path)
-                            screenshots.append(screenshot_path)
+                            # Find navigation buttons
+                            next_button = await page.query_selector('button[aria-label="Next page"]')
+                            prev_button = await page.query_selector('button[aria-label="Previous page"]')
+                            
+                            # If we're not on the first page, go back to first page
+                            if page_num == 0 and prev_button:
+                                for _ in range(50):  # Limit to avoid infinite loop
+                                    is_disabled = await prev_button.get_attribute('disabled')
+                                    if is_disabled:
+                                        break
+                                    await prev_button.click()
+                                    await page.wait_for_timeout(300)
+                            
+                            # Now navigate forward to desired page
+                            if page_num > 0 and next_button:
+                                for _ in range(page_num):
+                                    await next_button.click()
+                                    await page.wait_for_timeout(500)
+                            
+                            # Wait for the page content to load
+                            await page.wait_for_timeout(1000)
                         
-                        # Combine screenshots into PDF
-                        from reportlab.lib.pagesizes import letter
-                        from reportlab.pdfgen import canvas as pdf_canvas
-                        from PIL import Image
+                        # Wait for page to stabilize
+                        await page.wait_for_timeout(500)
+                        
+                        # Take the screenshot
+                        screenshot_path = os.path.join(temp_dir, f"page_{page_num + 1}.png")
+                        
+                        # Determine what to screenshot based on the viewer
+                        current_page_element = await page.evaluate("""
+                            () => {
+                                // First try getting the current visible page
+                                const pages = document.querySelectorAll('.drive-viewer-paginated-page');
+                                for (const page of pages) {
+                                    const rect = page.getBoundingClientRect();
+                                    if (rect.top < window.innerHeight && rect.bottom > 0) {
+                                        return {
+                                            x: Math.max(0, rect.left),
+                                            y: Math.max(0, rect.top),
+                                            width: Math.min(window.innerWidth, rect.width),
+                                            height: Math.min(window.innerHeight, rect.bottom - rect.top)
+                                        };
+                                    }
+                                }
+                                
+                                // Fallback: try to find the container
+                                const container = document.querySelector('.drive-viewer-paginated-scrollable');
+                                if (container) {
+                                    const rect = container.getBoundingClientRect();
+                                    return {
+                                        x: Math.max(0, rect.left),
+                                        y: Math.max(0, rect.top),
+                                        width: Math.min(window.innerWidth, rect.width),
+                                        height: Math.min(window.innerHeight, rect.bottom - rect.top)
+                                    };
+                                }
+                                
+                                // Last resort: screenshot the visible area
+                                return null;
+                            }
+                        """)
                         
-                        # Use the first image dimensions to determine page size
+                        if current_page_element:
+                            # Screenshot the specific page element
+                            await page.screenshot(path=screenshot_path, clip=current_page_element)
+                        else:
+                            # Screenshot the entire visible area
+                            await page.screenshot(path=screenshot_path)
+                        
+                        return screenshot_path
+                    
+                    # Capture all pages
+                    for i in range(page_count):
+                        logger.info(f"Capturing page {i+1} of {page_count}")
+                        screenshot_path = await capture_page(i)
+                        screenshots.append(screenshot_path)
+                        # Add progress indicator
+                        if (i+1) % 5 == 0 or i+1 == page_count:
+                            logger.info(f"Progress: {i+1}/{page_count} pages captured")
+                    
+                    # Combine screenshots into a PDF
+                    from PIL import Image
+                    from reportlab.lib.pagesizes import letter
+                    from reportlab.pdfgen import canvas as pdf_canvas
+                    
+                    logger.info(f"Combining {len(screenshots)} screenshots into PDF")
+                    
+                    # Use the first image dimensions to set PDF size if available
+                    if screenshots:
                         img = Image.open(screenshots[0])
                         img_width, img_height = img.size
-                        
                         c = pdf_canvas.Canvas(save_path, pagesize=(img_width, img_height))
                         
                         for screenshot in screenshots:
-                            img = Image.open(screenshot)
-                            c.drawImage(screenshot, 0, 0, img_width, img_height)
-                            c.showPage()
+                            # Check if file exists and has content
+                            if os.path.exists(screenshot) and os.path.getsize(screenshot) > 0:
+                                img = Image.open(screenshot)
+                                c.drawImage(screenshot, 0, 0, img_width, img_height)
+                                c.showPage()
                         
                         c.save()
-                        
-                        # Clean up
-                        for screenshot in screenshots:
+                    
+                    # Clean up screenshots
+                    for screenshot in screenshots:
+                        if os.path.exists(screenshot):
                             os.remove(screenshot)
-                        os.rmdir(temp_dir)
+                    os.rmdir(temp_dir)
+                    
+                    # Verify the PDF was created successfully
+                    if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
+                        logger.info(f"Successfully created PDF with {len(screenshots)} pages")
+                        return save_path
+                    else:
+                        logger.error("Failed to create PDF from screenshots")
+                        return None
                 else:
-                    # For other file types: Take a single screenshot
-                    temp_dir = tempfile.mkdtemp()
-                    screenshot_path = os.path.join(temp_dir, "screenshot.png")
+                    # For non-PDF files: take a single screenshot
+                    logger.info("Non-PDF file detected, taking single screenshot")
+                    screenshot_path = os.path.join(tempfile.gettempdir(), "screenshot.png")
                     await page.screenshot(path=screenshot_path, full_page=True)
                     
-                    # Determine final file type
-                    base, ext = os.path.splitext(save_path)
-                    
-                    if ext.lower() == '.pdf':
+                    # Convert to requested format if needed
+                    if save_path.lower().endswith('.pdf'):
                         # Convert to PDF
                         from PIL import Image
                         from reportlab.pdfgen import canvas as pdf_canvas
@@ -625,14 +801,13 @@ class DownloadManager:
                     
                     # Clean up
                     os.remove(screenshot_path)
-                    os.rmdir(temp_dir)
                 
                 # Close browser
                 await browser.close()
                 
                 # Verify file exists and is not empty
                 if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
-                    logger.info(f"Successfully downloaded view-only file to {save_path}")
+                    logger.info(f"Successfully downloaded file to {save_path}")
                     return save_path
                 else:
                     logger.error(f"Failed to create valid file at {save_path}")