import gradio as gr
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.firefox.service import Service as FirefoxService
from webdriver_manager.firefox import GeckoDriverManager
from PIL import Image
from io import BytesIO
import time

def take_screenshot_and_login(url, email, password):
    options = webdriver.FirefoxOptions()
    options.add_argument('--headless')
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')
    
    wd = None
    try:
        # Use GeckoDriverManager to automatically manage the Geckodriver executable
        service = FirefoxService(GeckoDriverManager().install())
        wd = webdriver.Firefox(service=service, options=options)
        
        wd.set_window_size(1080, 720)
        wd.get(url)
        
        # Log into Instagram
        email_input = wd.find_element(By.NAME, 'username')
        password_input = wd.find_element(By.NAME, 'password')
        
        email_input.send_keys(email)
        password_input.send_keys(password)
        password_input.send_keys(Keys.RETURN)
        
        # Wait for login to complete
        time.sleep(5)  # Adjust the sleep time if necessary
        
        # Take a screenshot
        screenshot = wd.get_screenshot_as_png()
    except Exception as e:
        print(f"Exception occurred: {e}")
        return Image.new('RGB', (1080, 720), color='white')
    finally:
        if wd:
            wd.quit()

    return Image.open(BytesIO(screenshot))

iface = gr.Interface(
    fn=take_screenshot_and_login,
    inputs=[
        gr.inputs.Textbox(label="Website URL", default="https://www.instagram.com/accounts/login/"),
        gr.inputs.Textbox(label="Email", default="example@gmail.com"),
        gr.inputs.Textbox(label="Password", default="example123", type="password")
    ],
    outputs=gr.outputs.Image(type="pil", label="Screenshot"),
    title="Instagram Login Screenshot Tool",
    description="Enter the URL of the Instagram login page, your email, and password to take a screenshot after logging in.",
)

if __name__ == "__main__":
    iface.launch()