Update app.py

app.py

@@ -43,6 +43,29 @@ def fetch_html(url: str):
     except Exception as e:
         return {"error": str(e)}
 
+@app.get("/download")
+def download_screenshot(url: str):
+    try:
+        with sync_playwright() as p:
+            browser = p.chromium.launch(headless=True)
+            page = browser.new_page()
+            
+            # Set custom headers before navigating to the page
+            headers = get_custom_headers(url)
+            page.set_extra_http_headers(headers)
+            
+            # Navigate to the URL and wait for network idle
+            page.goto(url, timeout=15000, wait_until="networkidle")
+            
+            # Take a screenshot of the page once content is fully loaded
+            screenshot = page.screenshot(full_page=True)  # Set full_page=True to capture the entire page
+            browser.close()
+        
+        # Return the screenshot as a response with image/png content type
+        return Response(content=screenshot, media_type="image/png")
+    except Exception as e:
+        return {"error": str(e)}
+
 if __name__ == "__main__":
     import uvicorn
     os.system("playwright install-deps")