Update app.py

app.py

@@ -3,36 +3,43 @@ import time
 import streamlit as st
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
+from selenium.webdriver.support.ui import WebDriverWait
+from selenium.webdriver.support import expected_conditions as EC
+from selenium.webdriver.common.by import By
 import tempfile
 
-# Function to capture a screenshot of a webpage
+# Function to capture a screenshot of a webpage after redirection
 def capture_screenshot(url, output_path):
     chrome_options = Options()
-    chrome_options.add_argument('--headless')  
-    chrome_options.add_argument('--disable-gpu')  
-    chrome_options.add_argument('--no-sandbox')  
-    chrome_options.add_argument('--window-size=1920,1080')  
+    chrome_options.add_argument('--headless')
+    chrome_options.add_argument('--disable-gpu')
+    chrome_options.add_argument('--no-sandbox')
+    chrome_options.add_argument('--window-size=1920,1080')
 
     driver = webdriver.Chrome(options=chrome_options)
     try:
         driver.get(url)
+
+        # Wait until the page is loaded and title stabilizes (handles JavaScript-based redirects)
+        WebDriverWait(driver, 15).until(lambda d: d.execute_script('return document.readyState') == 'complete')
+        time.sleep(2)  # Small extra wait to let JS redirects finish if any
+
         driver.save_screenshot(output_path)
     finally:
         driver.quit()
 
 # Streamlit UI
-st.title("Timed Web Screenshot Viewer")
+st.title("Timed Web Screenshot Viewer (with Redirect Support)")
 
 url = st.text_input("Enter a website URL:", "https://example.com")
 wait_time = st.number_input("Time to wait before taking screenshot (in seconds):", min_value=1, max_value=30, value=5, step=1)
 
 if st.button("Generate Screenshot"):
-    screenshot_path = None  # Initialize to ensure it's defined
+    screenshot_path = None
     try:
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
             screenshot_path = tmp_file.name
 
-        # Show a countdown
         with st.empty():
             for i in range(wait_time, 0, -1):
                 st.info(f"Waiting... {i} seconds left")