from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
from selenium_stealth import stealth
from webdriver_manager.chrome import ChromeDriverManager
import base64
import json
import requests
from dotenv import load_dotenv
import os
import time
import gradio as gr
import subprocess

def driverSetup():
    options = Options()
    options.add_argument("--no-sandbox")  # Required for Linux container environments
    options.add_argument("--disable-dev-shm-usage")  # Required for Linux container environments
    options.add_argument("--disable-gpu")  # Optional, can be enabled if needed
    options.add_argument("start-maximized")
    options.add_argument("--headless=new")
    options.add_experimental_option("excludeSwitches", ["enable-automation"])
    options.add_experimental_option('useAutomationExtension', False)
    options.add_argument("--window-size=1366,768")
    service = Service(ChromeDriverManager().install())
    driver = webdriver.Chrome(service=service, options=options)

    stealth(driver,
            languages=["en-US", "en"],
            vendor="Google Inc.",
            platform="Linux",
            webgl_vendor="MesaLib",
            renderer="Mesa DRI Intel(R) HD Graphics 5500",
            fix_hairline=True,
            )
    return driver

def screenshotName(url):
    if url.startswith('https://'):
        url = url.replace('https://', '')
    elif url.startswith('http://'):
        url = url.replace('http://', '')
    elif url.startswith('http://www.'):
        url = url.replace('http://www.', '')
    elif url.startswith('https://www.'):
        url = url.replace('https://www.', '')
    else:
        url = url
    if url.endswith('/'):
        url = url[:-1]
    else:
        url = url
    ssname = f"{url.replace('/', '')}.png"
    return ssname

def getScreenshots(driver, url):
    driver.get(url)
    ssname = screenshotName(url)
    driver.save_screenshot(ssname)
    return ssname

def saveScreenshot(url):
    driver = driverSetup()
    time.sleep(15)
    screenshot_file = getScreenshots(driver, url)
    driver.quit()
    return screenshot_file

def uploadandDelSS(file):
    load_dotenv()
    api = os.getenv("IMGBBAPI")
    url = os.getenv("IMGBBURL")
    filename = file
    with open(f"{filename}", "rb") as file:
        payload = {
            "key": api,
            "image": base64.b64encode(file.read()),
            "name": f"SS_{filename}",
            "expiration": "15552000"
        }
        r = requests.post(url, data= payload)
        view_url=(json.loads(r.text)["data"]["display_url"])
        file.close()
        os.remove(filename)
        return view_url, view_url

def main(url):
    ss = saveScreenshot(url)
    img, imgurl = uploadandDelSS(ss)
    return img, imgurl

app = gr.Interface(
    fn=main,
    inputs=[
        gr.Textbox(label="Enter URL", placeholder="https://google.com", type="text", interactive=True)
    ],
    outputs=[
        gr.Image(label="Website Screenshot"),
        gr.Textbox(label="Image URL", type="text", show_copy_button=True, interactive=False)
    ],
    title="Website Screenshot Capture<br> by <a href='https://nayankasturi.eu.org'>Nayan Kasturi</a> aka Raanna.<br> Checkout the <a href='https://github.com/raannakasturi'>Github</a> for more projects and contact info.",
    description="This app captures a screenshot of the website you enter and displays it.<br> Licenced under <a href='https://creativecommons.org/licenses/by-nc-sa/4.0/'>cc-by-nc-sa-4.0</a>",
    api_name="capture",
    concurrency_limit=10
)

def installChrome(): 
    # Update package list
    subprocess.run(['apt-get', 'update'])
    
    # Install wget and unzip
    subprocess.run(['apt-get', 'install', '-y', 'wget', 'unzip'])
    
    # Download Google Chrome
    subprocess.run(['wget', 'https://dl.google.com/linux/direct/google-chrome-stable_current_amd64.deb'])
    
    # Install Google Chrome
    subprocess.run(['apt-get', 'install', '-y', './google-chrome-stable_current_amd64.deb'])
    
    # Remove the downloaded package
    subprocess.run(['rm', 'google-chrome-stable_current_amd64.deb'])
    
    # Clean up
    subprocess.run(['apt-get', 'clean'])

if __name__ == "__main__":
    installChrome()
    print("App Starting...")
    app.launch()