import os
import streamlit as st
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from ascii_magic import AsciiArt

# Function to capture a screenshot of a webpage
def capture_screenshot(url, output_path='screenshot.png'):
    chrome_options = Options()
    chrome_options.add_argument('--headless')  # Run Chrome in headless mode
    chrome_options.add_argument('--disable-gpu')  # Disable GPU acceleration
    chrome_options.add_argument('--no-sandbox')  # Disable sandbox for better compatibility
    chrome_options.add_argument('--window-size=1920,1080')  # Set window size for consistent screenshots

    driver = webdriver.Chrome(options=chrome_options)
    try:
        driver.get(url)  # Open the URL
        driver.save_screenshot(output_path)  # Save screenshot to file
    finally:
        driver.quit()  # Ensure the driver is closed even if an error occurs

# Function to convert an image to ASCII art
def convert_to_ascii(image_path):
    try:
        # Correct usage of ascii_magic to convert image to ASCII
        ascii_art = AsciiArt.from_image(image_path, columns=100, mode='mono')
        return ascii_art
    except Exception as e:
        raise Exception(f"Failed to convert image to ASCII: {e}")

# Streamlit app layout
st.title("ASCII Web Screenshot Viewer")

# Input field for URL
url = st.text_input("Enter a website URL:", "https://example.com")

# Button to trigger the screenshot and conversion process
if st.button("Generate ASCII Screenshot"):
    try:
        screenshot_file = "screenshot.png"
        st.write("Capturing website screenshot...")
        capture_screenshot(url.strip('"'))  # Strip any extra quotes from the URL

        st.write("Converting to ASCII art...")
        ascii_output = convert_to_ascii(screenshot_file)  # Use the correct function

        # Display ASCII art in a text area
        st.text_area("ASCII Art Output", ascii_output, height=600)

    except Exception as e:
        st.error(f"An error occurred: {e}")
    finally:
        # Clean up the screenshot file if it exists
        if os.path.exists(screenshot_file):
            os.remove(screenshot_file)