import os
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 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')

    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 (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
    try:
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
            screenshot_path = tmp_file.name

        with st.empty():
            for i in range(wait_time, 0, -1):
                st.info(f"Waiting... {i} seconds left")
                time.sleep(1)

        st.write("Capturing screenshot...")
        capture_screenshot(url.strip('"'), screenshot_path)

        st.image(screenshot_path, caption='Screenshot of the webpage', use_column_width=True)

    except Exception as e:
        st.error(f"An error occurred: {e}")
    finally:
        if screenshot_path and os.path.exists(screenshot_path):
            os.remove(screenshot_path)