import gradio as gr
import asyncio
from playwright.async_api import async_playwright
import base64
from PIL import Image
import io

async def capture_screenshot(url):
    async with async_playwright() as p:
        browser = await p.chromium.launch()
        page = await browser.new_page()
        await page.goto(url)
        screenshot = await page.screenshot()
        await browser.close()
        return Image.open(io.BytesIO(screenshot))

async def open_link(url, index):
    try:
        screenshot = await capture_screenshot(url)
        return f"Request {index+1}: Successfully captured screenshot", screenshot
    except Exception as e:
        return f"Request {index+1}: Error: {str(e)}", None

async def open_links(link, count):
    tasks = [open_link(link, i) for i in range(count)]
    results = await asyncio.gather(*tasks)
    
    statuses = [result[0] for result in results]
    screenshots = [result[1] for result in results if result[1] is not None]
    
    return "\n".join(statuses), screenshots[0] if screenshots else None

def gradio_open_links(link, count):
    return asyncio.run(open_links(link, count))

iface = gr.Interface(
    fn=gradio_open_links,
    inputs=[
        gr.Textbox(label="Enter link", placeholder="https://example.com"),
        gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Number of times to open")
    ],
    outputs=[
        gr.Textbox(label="Results"),
        gr.Image(label="Screenshot")
    ],
    title="Link Opener with Screenshot",
    description="Enter a link and choose how many times to open it. The app will attempt to open the link, capture a screenshot, and display the first successful screenshot.",
    theme="huggingface",
    examples=[["https://www.example.com", 1]]
)

iface.launch()