from __future__ import annotations
from selenium import webdriver
from typing import Iterable

import gradio as Gradio
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes

import re
from PIL import Image
from io import BytesIO
from selenium.common.exceptions import WebDriverException

theme = Gradio.themes.Monochrome(
    primary_hue="purple",
    secondary_hue="purple",
    neutral_hue="neutral",
    radius_size=Gradio.themes.sizes.radius_sm,
    font=[Gradio.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)

class PurpleTheme(Base):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.purple,
        secondary_hue: colors.Color | str = colors.purple,
        neutral_hue: colors.Color | str = colors.neutral,
        spacing_size: sizes.Size | str = sizes.spacing_md,
        radius_size: sizes.Size | str = sizes.radius_md,
        font: fonts.Font
        | str
        | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Inter"),
            "ui-sans-serif",
            "sans-serif",
        ),
        font_mono: fonts.Font
        | str
        | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Space Grotesk"),
            "ui-monospace",
            "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            spacing_size=spacing_size,
            radius_size=radius_size,
            font=font,
            font_mono=font_mono,
        )
        super().set(
            button_primary_background_fill="linear-gradient(90deg, *primary_300, *secondary_400)",
            button_primary_background_fill_hover="linear-gradient(90deg, *primary_200, *secondary_300)",
            button_primary_text_color="white",
            button_primary_background_fill_dark="linear-gradient(90deg, *primary_600, *secondary_800)",
            block_shadow="*shadow_drop_lg",
            button_shadow="*shadow_drop_lg",
            input_background_fill="zinc",
            input_border_color="*secondary_300",
            input_shadow="*shadow_drop",
            input_shadow_focus="*shadow_drop_lg",
        )

custom_theme = PurpleTheme()

title = "Imaginor"
driver_type = 'chromedriver'
description = "<center>Imaginor is a small space for taking screenshots of websites. Handy for open-source projects that involves <br/>getting the site screenshots for status checks and view previews from a bird's eye view.</center>"

def run_imaginor(text: str):
    regex = r"^(https?://)"
    is_url = re.search(regex, text)

    if is_url:
        options = webdriver.ChromeOptions()
        options.add_argument('--headless')
        options.add_argument('--no-sandbox')
        options.add_argument('--disable-dev-shm-usage')

        try:
            driver = webdriver.Chrome(options=options)
            driver.get(text)
            driver.implicitly_wait(30)
            driver.set_window_size(1920, 1080)
            screenshot = driver.get_screenshot_as_png()
        except WebDriverException as e:
            return [Image.new('RGB', (1, 1)), 'Website imagined by Imaginor, operation failed.']
        finally:
            if driver:
                driver.quit()

        return [Image.open(BytesIO(screenshot)), 'Website imagined by Imaginor, operation success.']
    else:
        return [None, 'Please enter a valid URL of a website/host.']

Gradio.Interface(
    run_imaginor,
    Gradio.Textbox(placeholder="Enter your URL here", label="Website URL / Endpoint"),
    [Gradio.Image(type="pil", label="Screenshot"),
    Gradio.Textbox(label="Text result", interactive=False)],
    title=title,
    description=description,
    theme=custom_theme,
    analytics_enabled=False,
    css=".generating {visibility: hidden}"
).launch(debug=True)