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.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.options import Options
import time
import subprocess
import os

# Function to log Chrome and ChromeDriver versions
def log_versions():
    try:
        # Print PATH for debugging
        print(f"System PATH: {os.environ.get('PATH')}")
        # Try finding Chrome binary
        chrome_binary = None
        for binary in ["google-chrome", "google-chrome-stable", "chromium-browser"]:
            try:
                chrome_path = subprocess.check_output(["which", binary]).decode().strip()
                print(f"Found {binary} at: {chrome_path}")
                chrome_binary = binary
                break
            except subprocess.CalledProcessError:
                continue
        if not chrome_binary:
            print("Error: No Chrome binary (google-chrome, google-chrome-stable, chromium-browser) found")
            return

        # Get Chrome version
        chrome_version = subprocess.check_output([chrome_binary, "--version"]).decode().strip()
        print(f"Chrome version: {chrome_version}")

        # Try finding ChromeDriver
        try:
            chromedriver_path = subprocess.check_output(["which", "chromedriver"]).decode().strip()
            print(f"Found chromedriver at: {chromedriver_path}")
            chromedriver_version = subprocess.check_output(["chromedriver", "--version"]).decode().strip()
            print(f"ChromeDriver version: {chromedriver_version}")
        except subprocess.CalledProcessError as e:
            print(f"Error finding chromedriver: {str(e)}")

    except Exception as e:
        print(f"Error checking versions: {str(e)}")

# Function to generate image URL
def generate_image(prompt):
    # Log versions for debugging
    log_versions()

    # Configure Chrome options for headless mode
    chrome_options = Options()
    chrome_options.add_argument("--headless")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("window-size=1920x1080")

    # Initialize the WebDriver
    try:
        driver = webdriver.Chrome(options=chrome_options)
    except Exception as e:
        return f"Error initializing ChromeDriver: {str(e)}"

    try:
        # Open the website
        driver.get("https://gptimage.ai/")

        # Wait for the textarea input field
        input_field = WebDriverWait(driver, 10).until(
            EC.presence_of_element_located(
                (By.CSS_SELECTOR, "textarea.wp-content-visualizer-chat-input")
            )
        )
        # Enter the prompt
        input_field.send_keys(prompt)

        # Wait for the "Generate Image" button and click it
        generate_button = WebDriverWait(driver, 10).until(
            EC.element_to_be_clickable(
                (By.CSS_SELECTOR, "button.wp-content-visualizer-chat-send")
            )
        )
        generate_button.click()

        # Wait for the generated image
        generated_image = WebDriverWait(driver, 20).until(
            EC.presence_of_element_located(
                (By.CSS_SELECTOR, "img.wp-content-visualizer-generated-image")
            )
        )

        # Extract the image URL
        image_url = generated_image.get_attribute("src")
        return image_url

    except Exception as e:
        return f"Error: {str(e)}"
    finally:
        # Close the browser
        driver.quit()

# Gradio interface
def gradio_interface(prompt):
    try:
        image_url = generate_image(prompt)
        if image_url.startswith("Error"):
            return image_url, None
        return f"Generated Image URL: {image_url}", image_url
    except Exception as e:
        return f"Error: {str(e)}", None

# Create Gradio app
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g., A serene mountain landscape at sunset"),
    outputs=[gr.Textbox(label="Result"), gr.Image(label="Generated Image")],
    title="Image Generator with selenium",
    description="Enter a prompt to generate an image using gptimage.ai."
)

# Launch the app
if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)  # No share=True