capture_screenshots.py

"""
Screenshot capture tool for ChatUI Helper documentation
Uses Playwright to automate browser interactions and capture screenshots
"""

import asyncio
import os
from pathlib import Path
from playwright.async_api import async_playwright
import logging
from datetime import datetime

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
logger = logging.getLogger(__name__)

# Screenshot configurations for each documentation step
SCREENSHOT_CONFIGS = [
    {
        "name": "main_interface",
        "url": "http://localhost:7860",
        "description": "Main interface with all tabs",
        "actions": [],
        "wait_time": 2000,
        "selector": None,  # Full page
        "output": "img/main_interface.png"
    },
    {
        "name": "config_tab",
        "url": "http://localhost:7860",
        "description": "Configuration tab",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"}
        ],
        "wait_time": 1000,
        "selector": None,
        "output": "img/config_tab.png"
    },
    {
        "name": "template_dropdown",
        "url": "http://localhost:7860",
        "description": "Template selection dropdown",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"},
            {"type": "click", "selector": "div[data-testid='dropdown']", "nth": 0}
        ],
        "wait_time": 1000,
        "selector": "div.dropdown-menu",
        "output": "img/template_dropdown.png"
    },
    {
        "name": "preview_tab",
        "url": "http://localhost:7860",
        "description": "Preview tab with chat interface",
        "actions": [
            {"type": "click", "selector": "button:has-text('Preview')"}
        ],
        "wait_time": 2000,
        "selector": None,
        "output": "img/preview_tab.png"
    },
    {
        "name": "preview_chat_example",
        "url": "http://localhost:7860",
        "description": "Preview tab with example conversation",
        "actions": [
            {"type": "click", "selector": "button:has-text('Preview')"},
            {"type": "wait", "time": 2000},
            {"type": "type", "selector": "textarea[placeholder*='Type a message']", "text": "Hello! Can you help me understand quantum mechanics?"},
            {"type": "click", "selector": "button:has-text('Send')"},
            {"type": "wait", "time": 3000}
        ],
        "wait_time": 1000,
        "selector": None,
        "output": "img/preview_chat_example.png"
    },
    {
        "name": "docs_tab",
        "url": "http://localhost:7860",
        "description": "Documentation tab",
        "actions": [
            {"type": "click", "selector": "button:has-text('Docs')"}
        ],
        "wait_time": 1000,
        "selector": None,
        "output": "img/docs_tab.png"
    },
    {
        "name": "docs_quickstart",
        "url": "http://localhost:7860",
        "description": "Documentation quick start guide",
        "actions": [
            {"type": "click", "selector": "button:has-text('Docs')"},
            {"type": "wait", "time": 1000},
            {"type": "click", "selector": "button:has-text('📖 Quick Start Guide')"}
        ],
        "wait_time": 1000,
        "selector": None,
        "output": "img/docs_quickstart.png"
    },
    {
        "name": "config_system_prompt",
        "url": "http://localhost:7860",
        "description": "System configuration section",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"},
            {"type": "wait", "time": 1000},
            {"type": "scroll", "y": 500}
        ],
        "wait_time": 1000,
        "selector": None,
        "output": "img/config_system_prompt.png"
    },
    {
        "name": "config_api_section",
        "url": "http://localhost:7860",
        "description": "API configuration section",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"},
            {"type": "wait", "time": 1000},
            {"type": "scroll", "y": 1200}
        ],
        "wait_time": 1000,
        "selector": None,
        "output": "img/config_api_section.png"
    },
    {
        "name": "generate_button",
        "url": "http://localhost:7860",
        "description": "Generate deployment package button",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"},
            {"type": "wait", "time": 1000},
            {"type": "scroll", "y": 2000}
        ],
        "wait_time": 1000,
        "selector": "button:has-text('🗳️ Generate Deployment Package')",
        "output": "img/generate_button.png"
    }
]

async def capture_screenshot(page, config):
    """Capture a single screenshot based on configuration"""
    logger.info(f"Capturing screenshot: {config['name']} - {config['description']}")
    
    try:
        # Navigate to URL
        await page.goto(config['url'])
        
        # Execute actions
        for action in config.get('actions', []):
            if action['type'] == 'click':
                selector = action['selector']
                nth = action.get('nth', 0)
                if nth > 0:
                    await page.locator(selector).nth(nth).click()
                else:
                    await page.click(selector)
                await page.wait_for_timeout(500)
                
            elif action['type'] == 'type':
                await page.fill(action['selector'], action['text'])
                
            elif action['type'] == 'wait':
                await page.wait_for_timeout(action['time'])
                
            elif action['type'] == 'scroll':
                await page.evaluate(f"window.scrollBy(0, {action['y']})")
                await page.wait_for_timeout(500)
        
        # Wait for any animations to complete
        await page.wait_for_timeout(config.get('wait_time', 1000))
        
        # Take screenshot
        output_path = Path(config['output'])
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        if config.get('selector'):
            # Screenshot of specific element
            element = page.locator(config['selector'])
            await element.screenshot(path=str(output_path))
        else:
            # Full page screenshot
            await page.screenshot(path=str(output_path), full_page=False)
            
        logger.info(f"✅ Saved screenshot to {output_path}")
        return True
        
    except Exception as e:
        logger.error(f"❌ Failed to capture {config['name']}: {str(e)}")
        return False

async def capture_all_screenshots():
    """Capture all configured screenshots"""
    logger.info("Starting screenshot capture process...")
    
    # Ensure the app is running
    logger.info("Make sure the Gradio app is running on http://localhost:7860")
    
    async with async_playwright() as p:
        # Launch browser
        browser = await p.chromium.launch(headless=False)  # Set to True for headless mode
        context = await browser.new_context(
            viewport={'width': 1280, 'height': 800},
            device_scale_factor=2  # Higher quality screenshots
        )
        page = await context.new_page()
        
        # Capture each screenshot
        success_count = 0
        for config in SCREENSHOT_CONFIGS:
            if await capture_screenshot(page, config):
                success_count += 1
            await page.wait_for_timeout(1000)  # Brief pause between screenshots
        
        # Clean up
        await browser.close()
        
    logger.info(f"\nCapture complete! {success_count}/{len(SCREENSHOT_CONFIGS)} screenshots captured successfully.")
    
    # Create a summary report
    report_path = Path("img/screenshot_report.txt")
    with open(report_path, 'w') as f:
        f.write(f"Screenshot Capture Report\n")
        f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Total Screenshots: {len(SCREENSHOT_CONFIGS)}\n")
        f.write(f"Successful: {success_count}\n\n")
        
        f.write("Screenshots Generated:\n")
        for config in SCREENSHOT_CONFIGS:
            f.write(f"- {config['output']}: {config['description']}\n")
    
    logger.info(f"Report saved to {report_path}")

async def capture_single_screenshot(name):
    """Capture a single screenshot by name"""
    config = next((c for c in SCREENSHOT_CONFIGS if c['name'] == name), None)
    if not config:
        logger.error(f"No configuration found for screenshot: {name}")
        return
    
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=False)
        context = await browser.new_context(
            viewport={'width': 1280, 'height': 800},
            device_scale_factor=2
        )
        page = await context.new_page()
        
        await capture_screenshot(page, config)
        
        await browser.close()

def main():
    """Main entry point"""
    import sys
    
    if len(sys.argv) > 1:
        # Capture specific screenshot
        asyncio.run(capture_single_screenshot(sys.argv[1]))
    else:
        # Capture all screenshots
        asyncio.run(capture_all_screenshots())

if __name__ == "__main__":
    main()