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

def capture_screenshot(url, output_path='screenshot.png'):
    chrome_options = Options()
    chrome_options.add_argument('--headless')
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--window-size=1920,1080')

    driver = webdriver.Chrome(options=chrome_options)
    driver.get(url)
    driver.save_screenshot(output_path)
    driver.quit()

def convert_to_ascii(image_path):
    output = ascii_magic.from_image_path(image_path, columns=100, mode=ascii_magic.Modes.ASCII)
    return str(output)

st.title("ASCII Web Screenshot Viewer")

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

if st.button("Generate ASCII Screenshot"):
    try:
        screenshot_file = "screenshot.png"
        st.write("Capturing website screenshot...")
        capture_screenshot(url, screenshot_file)

        st.write("Converting to ASCII art...")
        ascii_output = convert_to_ascii(screenshot_file)

        st.text_area("ASCII Art Output", ascii_output, height=600)
        os.remove(screenshot_file)

    except Exception as e:
        st.error(f"An error occurred: {e}")