capture_docs_screenshots.py

"""
Specialized screenshot capture for documentation steps in support_docs.py
Captures specific HuggingFace Spaces deployment screenshots
"""

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

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

# Documentation-specific screenshot configurations
DOCS_SCREENSHOTS = {
    # Step 1: Generate & Create Space
    "huggingface_spaces_page": {
        "url": "https://huggingface.co/spaces",
        "description": "HuggingFace Spaces main page",
        "output": "img/hf_spaces_main.png",
        "actions": [],
        "wait_time": 3000
    },
    "new_space_button": {
        "url": "https://huggingface.co/spaces",
        "description": "New Space button",
        "output": "img/hf_new_space_button.png",
        "actions": [],
        "selector": "button:has-text('New Space')",
        "wait_time": 2000
    },
    "space_creation_form": {
        "url": "https://huggingface.co/new-space",
        "description": "Space creation form (img17.png replacement)",
        "output": "img/img17.png",
        "actions": [],
        "wait_time": 3000
    },
    "hardware_selection": {
        "url": "https://huggingface.co/new-space",
        "description": "Hardware selection options (img16.png replacement)",
        "output": "img/img16.png",
        "actions": [
            {"type": "scroll", "y": 400}
        ],
        "wait_time": 2000
    },
    
    # Step 2: Upload Files
    "files_tab": {
        "url": "https://huggingface.co/spaces/{username}/{space_name}/tree/main",
        "description": "Files tab in Space",
        "output": "img/hf_files_tab.png",
        "actions": [],
        "wait_time": 2000
    },
    "upload_files_button": {
        "url": "https://huggingface.co/spaces/{username}/{space_name}/tree/main",
        "description": "Upload files interface (img12.png replacement)",
        "output": "img/img12.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('Upload files')"}
        ],
        "wait_time": 2000
    },
    "files_after_upload": {
        "url": "https://huggingface.co/spaces/{username}/{space_name}/tree/main",
        "description": "Files view after upload (img8.png replacement)",
        "output": "img/img8.png",
        "actions": [],
        "wait_time": 2000
    },
    
    # Step 3: Configure API Secrets
    "settings_tab": {
        "url": "https://huggingface.co/spaces/{username}/{space_name}/settings",
        "description": "Settings tab navigation (img4.png replacement)",
        "output": "img/img4.png",
        "actions": [],
        "wait_time": 2000
    },
    "variables_secrets_section": {
        "url": "https://huggingface.co/spaces/{username}/{space_name}/settings",
        "description": "Variables and secrets section",
        "output": "img/hf_variables_secrets.png",
        "actions": [
            {"type": "scroll", "y": 300}
        ],
        "wait_time": 2000
    },
    "new_secret_form": {
        "url": "https://huggingface.co/spaces/{username}/{space_name}/settings",
        "description": "New secret configuration form (img3.png replacement)",
        "output": "img/img3.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('New secret')"}
        ],
        "wait_time": 2000
    },
    
    # Step 4: Monitor Build
    "build_logs": {
        "url": "https://huggingface.co/spaces/{username}/{space_name}",
        "description": "Build process logs (img7.png replacement)",
        "output": "img/img7.png",
        "actions": [],
        "wait_time": 3000
    },
    "successful_deployment": {
        "url": "https://huggingface.co/spaces/{username}/{space_name}",
        "description": "Successfully deployed Space (img1.png replacement)",
        "output": "img/img1.png",
        "actions": [],
        "wait_time": 5000
    }
}

# Local app screenshots for main documentation
LOCAL_APP_SCREENSHOTS = {
    "config_tab_full": {
        "url": "http://localhost:7860",
        "description": "Full configuration tab",
        "output": "img/config_tab_full.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"}
        ],
        "wait_time": 2000,
        "full_page": True
    },
    "template_selection": {
        "url": "http://localhost:7860",
        "description": "Template selection with dropdown open",
        "output": "img/template_selection.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"},
            {"type": "wait", "time": 1000},
            {"type": "click", "selector": "div.gr-dropdown"}
        ],
        "wait_time": 1500
    },
    "space_identity_section": {
        "url": "http://localhost:7860",
        "description": "Space identity configuration",
        "output": "img/space_identity_config.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"},
            {"type": "scroll", "y": 200}
        ],
        "wait_time": 1000
    },
    "system_config_section": {
        "url": "http://localhost:7860",
        "description": "System configuration section",
        "output": "img/system_config_section.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"},
            {"type": "scroll", "y": 600}
        ],
        "wait_time": 1000
    },
    "api_config_section": {
        "url": "http://localhost:7860",
        "description": "API configuration section",
        "output": "img/api_config_section.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('Configure Space')"},
            {"type": "scroll", "y": 1200}
        ],
        "wait_time": 1000
    },
    "preview_tab_active": {
        "url": "http://localhost:7860",
        "description": "Preview tab with active chat",
        "output": "img/preview_tab_active.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('Preview')"},
            {"type": "wait", "time": 2000},
            {"type": "type", "selector": "textarea", "text": "Tell me about quantum mechanics"},
            {"type": "click", "selector": "button:has-text('Send')"},
            {"type": "wait", "time": 4000}
        ],
        "wait_time": 1000
    },
    "docs_tab_overview": {
        "url": "http://localhost:7860",
        "description": "Documentation tab overview",
        "output": "img/docs_tab_overview.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('Docs')"}
        ],
        "wait_time": 2000
    },
    "docs_deployment_steps": {
        "url": "http://localhost:7860",
        "description": "Deployment steps accordion expanded",
        "output": "img/docs_deployment_steps.png",
        "actions": [
            {"type": "click", "selector": "button:has-text('Docs')"},
            {"type": "wait", "time": 1000},
            {"type": "click", "selector": "button:has-text('🗳️ Step 3: Generate & Deploy')"}
        ],
        "wait_time": 2000
    }
}

async def capture_screenshot(page, config, replace_values=None):
    """Capture a single screenshot based on configuration"""
    url = config['url']
    if replace_values:
        for key, value in replace_values.items():
            url = url.replace(f"{{{key}}}", value)
    
    logger.info(f"Capturing: {config['description']} -> {config['output']}")
    
    try:
        # Navigate to URL
        await page.goto(url, wait_until='networkidle')
        
        # Execute actions
        for action in config.get('actions', []):
            if action['type'] == 'click':
                await page.click(action['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
        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)
        
        screenshot_options = {
            "path": str(output_path),
            "full_page": config.get('full_page', False)
        }
        
        if config.get('selector'):
            element = page.locator(config['selector'])
            await element.screenshot(path=str(output_path))
        else:
            await page.screenshot(**screenshot_options)
            
        logger.info(f"✅ Saved: {output_path}")
        return True
        
    except Exception as e:
        logger.error(f"❌ Failed: {config['description']} - {str(e)}")
        return False

async def capture_huggingface_screenshots(username=None, space_name=None):
    """Capture HuggingFace-specific screenshots"""
    logger.info("Capturing HuggingFace screenshots...")
    
    if not username or not space_name:
        logger.warning("No username/space_name provided. Using placeholders.")
        username = "your-username"
        space_name = "your-space-name"
    
    replace_values = {
        "username": username,
        "space_name": space_name
    }
    
    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()
        
        for name, config in DOCS_SCREENSHOTS.items():
            await capture_screenshot(page, config, replace_values)
            await page.wait_for_timeout(1000)
        
        await browser.close()

async def capture_local_app_screenshots():
    """Capture local app screenshots"""
    logger.info("Capturing local app screenshots...")
    logger.info("Make sure the app is running on http://localhost:7860")
    
    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()
        
        for name, config in LOCAL_APP_SCREENSHOTS.items():
            await capture_screenshot(page, config)
            await page.wait_for_timeout(1000)
        
        await browser.close()

async def main():
    """Main entry point"""
    import sys
    
    if len(sys.argv) > 1:
        if sys.argv[1] == "local":
            await capture_local_app_screenshots()
        elif sys.argv[1] == "huggingface":
            username = sys.argv[2] if len(sys.argv) > 2 else None
            space_name = sys.argv[3] if len(sys.argv) > 3 else None
            await capture_huggingface_screenshots(username, space_name)
        else:
            logger.error("Usage: python capture_docs_screenshots.py [local|huggingface] [username] [space_name]")
    else:
        # Capture both sets
        await capture_local_app_screenshots()
        logger.info("\nNow capturing HuggingFace screenshots...")
        await capture_huggingface_screenshots()

if __name__ == "__main__":
    asyncio.run(main())