Added middleware

webrify2.py

@@ -186,11 +186,63 @@ async def get_metadata(url: str):
 #     finally:
 #         await browser.close()
 #         await pw.stop()
+# @app.get("/screenshot", response_model=ScreenshotResponse)
+# async def get_screenshot(url: str):
+#     page, browser, pw = await get_page(url)
+#     try:
+#         # Scroll to bottom to trigger lazy-loaded content
+#         await page.evaluate("""
+#             () => {
+#                 return new Promise((resolve) => {
+#                     let totalHeight = 0;
+#                     const distance = 100;
+#                     const timer = setInterval(() => {
+#                         window.scrollBy(0, distance);
+#                         totalHeight += distance;
+#                         if (totalHeight >= document.body.scrollHeight) {
+#                             clearInterval(timer);
+#                             resolve();
+#                         }
+#                     }, 100);
+#                 });
+#             }
+#         """)
+
+#         # Give time for images and content to load
+#         await page.wait_for_timeout(2000)
+
+#         image_bytes = await page.screenshot(full_page=True)
+#         image_base64 = base64.b64encode(image_bytes).decode()
+#         return {"screenshot": image_base64}
+#     finally:
+#         await browser.close()
+#         await pw.stop()
+
 @app.get("/screenshot", response_model=ScreenshotResponse)
 async def get_screenshot(url: str):
     page, browser, pw = await get_page(url)
     try:
-        # Scroll to bottom to trigger lazy-loaded content
+        # Go to the page and wait until the network is idle
+        await page.goto(url, wait_until="networkidle", timeout=60000)
+
+        # Wait for the header (or similar element) to load
+        try:
+            await page.wait_for_selector("header", timeout=10000)
+        except:
+            pass  # Don't fail if the header doesn't exist
+
+        # Remove sticky or fixed header issues before full-page screenshot
+        await page.add_style_tag(content="""
+            * {
+                scroll-behavior: auto !important;
+            }
+            header, .sticky, .fixed, [style*="position:fixed"] {
+                position: static !important;
+                top: auto !important;
+            }
+        """)
+
+        # Scroll down to trigger lazy loading
         await page.evaluate("""
             () => {
                 return new Promise((resolve) => {
@@ -208,11 +260,13 @@ async def get_screenshot(url: str):
             }
         """)
 
-        # Give time for images and content to load
+        # Wait to ensure lazy content and animations complete
         await page.wait_for_timeout(2000)
 
+        # Take full-page screenshot
         image_bytes = await page.screenshot(full_page=True)
         image_base64 = base64.b64encode(image_bytes).decode()
+
         return {"screenshot": image_base64}
     finally:
         await browser.close()