import os, tempfile, time
import gradio as gr
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from PIL import Image
from io import BytesIO

def html_to_screenshot(html_code: str) -> Image.Image:
    # Configure Selenium to use headless Chrome
    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")

    # Launch headless Chrome
    driver = webdriver.Chrome(options=chrome_options)
    try:
        # Write HTML code to a temporary file
        tmp_file = tempfile.NamedTemporaryFile(suffix=".html", delete=False)
        tmp_path = tmp_file.name
        tmp_file.write(html_code.encode('utf-8'))
        tmp_file.close()
        driver.get(f"file://{tmp_path}")
        time.sleep(1)  # allow any dynamic content to load if needed

        # Get total page dimensions
        total_width = driver.execute_script(
            "return Math.max(document.body.scrollWidth, document.documentElement.scrollWidth);")
        total_height = driver.execute_script(
            "return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight);")
        # Set the browser window to full content width and a fixed viewport height
        viewport_height = 1000  # px
        driver.set_window_size(total_width, viewport_height)

        # Mark all fixed or sticky elements
        driver.execute_script(
            "document.querySelectorAll('*').forEach(el => {"
            "  const pos = window.getComputedStyle(el).position;"
            "  if(pos === 'fixed' || pos === 'sticky') { el.setAttribute('data-fixed', 'true'); }"
            "});")
        # Screenshot the top of the page (with sticky elements visible)
        screenshots = []
        png_data = driver.get_screenshot_as_png()
        screenshots.append(Image.open(BytesIO(png_data)))

        # Hide sticky/fixed elements before taking further screenshots
        driver.execute_script(
            "document.querySelectorAll('[data-fixed=\"true\"]').forEach(el => el.style.visibility='hidden');")

        # Scroll and capture screenshots until reaching the bottom
        pixels_scrolled = viewport_height
        while True:
            if pixels_scrolled >= total_height:
                break  # done if we've covered the whole height
            driver.execute_script(f"window.scrollTo(0, {pixels_scrolled});")
            time.sleep(0.2)
            # Check actual scroll position in case we hit the bottom
            current_offset = driver.execute_script("return window.pageYOffset;")
            if current_offset < pixels_scrolled:
                # Adjust if we couldn't scroll the full amount (at bottom of page)
                current_offset = total_height - viewport_height
                driver.execute_script(f"window.scrollTo(0, {current_offset});")
                time.sleep(0.1)
            # Capture screenshot at the current offset
            png_data = driver.get_screenshot_as_png()
            screenshots.append(Image.open(BytesIO(png_data)))
            # Prepare for next iteration
            pixels_scrolled = current_offset + viewport_height

        # Stitch screenshots into one tall image
        # Compute overlap if the last screenshot went beyond the content bottom
        remainder = total_height % viewport_height
        overlap = viewport_height - remainder if remainder != 0 else 0
        if overlap and len(screenshots) > 1:
            # Crop the overlapping top part from the last image
            last_img = screenshots[-1]
            screenshots[-1] = last_img.crop((0, overlap, last_img.width, last_img.height))
        # Combine images vertically
        total_combined_height = sum(img.height for img in screenshots)
        combined_img = Image.new("RGB", (total_width, total_combined_height))
        y = 0
        for img in screenshots:
            combined_img.paste(img, (0, y))
            y += img.height

        return combined_img
    finally:
        driver.quit()

# Set up Gradio interface
interface = gr.Interface(
    fn=html_to_screenshot,
    inputs=gr.Textbox(label="HTML Code", lines=15),
    outputs=gr.Image(type="pil"),
    title="HTML Full-Page Screenshot",
    description="Enter HTML code and generate a full-page screenshot image."
)

if __name__ == "__main__":
    interface.launch()