import gradio as gr
import imgkit
from PIL import Image
from io import BytesIO

def html_to_image(html_code, width, height, zoom, jsdelay):
    try:
        options = {
            'format': 'png',
            'encoding': 'UTF-8',
            'width': str(width),
            'height': str(height),
            'zoom': str(zoom),
            'javascript-delay': str(jsdelay),
        }
        image = Image.open(BytesIO(imgkit.from_string(html_code, False, options=options)))
        return image
    except Exception as e:
        raise gr.Error(e, duration=5)

def website_to_image(website_link, width, height, zoom, jsdelay):
    try:
        if not website_link:
            website_link = "https://www.google.com"
        options = {
            'format': 'png',
            'encoding': 'UTF-8',
            'width': str(width),
            'height': str(height),
            'zoom': str(zoom),
            'javascript-delay': str(jsdelay),
        }
        image = Image.open(BytesIO(imgkit.from_url(website_link, False, options=options)))
        return image
    except Exception as e:
        raise gr.Error(e, duration=5)

with gr.Blocks() as interface:
    gr.Markdown("# HTML to Image Converter")
    gr.Markdown("Enter HTML code and set dimensions to generate an image")
    with gr.Row():
        with gr.Column():
            html_code_input = gr.Code(
                label="HTML Code",
                language="html",
                lines=30,
                interactive=True
            )
            width_input = gr.Number(
                label="Width",
                value=1280,
                step=10
            )
            height_input = gr.Number(
                label="Height",
                value=720,
                step=10,
                )
            zoom_input = gr.Number(
                label="Zoom",
                value=1.0,
                step=0.1
            )
            jsdelay_input = gr.Number(
                label="Javascript Delay",
                value=200,
                step=100
            )
            with gr.Row(variant="panel", ):
                with gr.Column():
                    website_link = gr.Textbox(
                        label="Website Link",
                        placeholder="https://www.google.com"
                    )
                snapshotwebsite = gr.Button("Show Website", size="lg")
        output_image = gr.Image(type="pil", format="PNG", show_label=False)
    generate_button = gr.Button("Force Refresh")
    
    html_code_input.change(
        fn=html_to_image,
        inputs=[html_code_input, width_input, height_input, zoom_input, jsdelay_input],
        outputs=output_image
    )
    width_input.change(
        fn=html_to_image,
        inputs=[html_code_input, width_input, height_input, zoom_input, jsdelay_input],
        outputs=output_image
    )
    height_input.change(
        fn=html_to_image,
        inputs=[html_code_input, width_input, height_input, zoom_input, jsdelay_input],
        outputs=output_image
    )
    zoom_input.change(
        fn=html_to_image,
        inputs=[html_code_input, width_input, height_input, zoom_input, jsdelay_input],
        outputs=output_image
    )
    
    jsdelay_input.change(
        fn=html_to_image,
        inputs=[html_code_input, width_input, height_input, zoom_input, jsdelay_input],
        outputs=output_image
    )

    generate_button.click(
        fn=html_to_image,
        inputs=[html_code_input, width_input, height_input, zoom_input, jsdelay_input],
        outputs=output_image
    )
    
    snapshotwebsite.click(
        fn=website_to_image,
        inputs=[website_link, width_input, height_input, zoom_input, jsdelay_input],
        outputs=output_image
    )

    gr.Examples(
        examples=[
            ["<div style='background: linear-gradient(45deg, #ff6b6b, #4ecdc4); padding: 20px; border-radius: 10px;'><h1 style='color: white; font-family: Arial;'>Hello, World!</h1></div>", 800, 400, 1.0, 200],
            ["<div style='background: #f0f0f0; padding: 20px;'><ul style='color: #333;'><li>Item 1</li><li>Item 2</li><li>Item 3</li></ul></div>", 600, 300, 1.0, 200],
        ],
        inputs=[html_code_input, width_input, height_input, zoom_input, jsdelay_input],
    )

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