app.py

import asyncio
import base64
import os
import re
from urllib.parse import urlparse

from flask import Flask, request, jsonify, send_file, render_template, Response
from playwright.async_api import async_playwright, TimeoutError as PlaywrightTimeoutError

app = Flask(__name__)

# Global browser instance (lazy initialisation)
_browser = None
_playwright = None


async def get_browser():
    """Get or create a Playwright browser instance."""
    global _browser, _playwright
    if _browser is None:
        _playwright = await async_playwright().start()
        _browser = await _playwright.chromium.launch(
            headless=True,
            args=['--no-sandbox', '--disable-setuid-sandbox']
        )
    return _browser


@app.route('/')
def index():
    """Serve the main page."""
    return render_template('index.html')


@app.route('/preview', methods=['POST'])
async def preview():
    """
    Generate a screenshot thumbnail of the given URL.
    Expects JSON: { "url": "...", "wait_time": 2000 (optional) }
    Returns base64-encoded image.
    """
    data = request.get_json()
    if not data or 'url' not in data:
        return jsonify({'error': 'URL is required'}), 400

    url = data['url'].strip()
    if not is_valid_url(url):
        return jsonify({'error': 'Invalid URL'}), 400

    wait_time = data.get('wait_time', 2000)  # default 2 seconds

    try:
        browser = await get_browser()
        context = await browser.new_context(
            viewport={'width': 1280, 'height': 720},
            user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        )
        page = await context.new_page()

        # Navigate with timeout
        await page.goto(url, wait_until='networkidle', timeout=30000)
        # Wait additional time for JS execution
        await page.wait_for_timeout(wait_time)

        # Take a full-page screenshot
        screenshot = await page.screenshot(full_page=True, type='jpeg', quality=80)
        await context.close()

        # Encode to base64
        img_base64 = base64.b64encode(screenshot).decode('utf-8')
        return jsonify({'image': f'data:image/jpeg;base64,{img_base64}'})

    except PlaywrightTimeoutError:
        return jsonify({'error': 'Timeout: The page did not load in time.'}), 408
    except Exception as e:
        # Log error (consider adding proper logging)
        return jsonify({'error': f'Preview failed: {str(e)}'}), 500


@app.route('/generate', methods=['POST'])
async def generate_pdf():
    """
    Generate a PDF from the given URL with options.
    Expects JSON: {
        "url": "...",
        "page_size": "A4" | "Letter",
        "margin": "0.5" (in inches),
        "dark_mode": true/false,
        "wait_time": 2000 (ms)
    }
    Returns the PDF file as attachment.
    """
    data = request.get_json()
    if not data or 'url' not in data:
        return jsonify({'error': 'URL is required'}), 400

    url = data['url'].strip()
    if not is_valid_url(url):
        return jsonify({'error': 'Invalid URL'}), 400

    page_size = data.get('page_size', 'A4')
    margin = data.get('margin', 0.5)
    dark_mode = data.get('dark_mode', False)
    wait_time = data.get('wait_time', 2000)

    try:
        browser = await get_browser()
        context = await browser.new_context(
            viewport={'width': 1280, 'height': 720},
            user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        )
        page = await context.new_page()

        # Set dark mode if requested
        if dark_mode:
            await page.emulate_media(color_scheme='dark')

        await page.goto(url, wait_until='networkidle', timeout=30000)
        await page.wait_for_timeout(wait_time)

        # Generate PDF
        pdf_bytes = await page.pdf(
            format=page_size,
            margin={
                'top': f'{margin}in',
                'bottom': f'{margin}in',
                'left': f'{margin}in',
                'right': f'{margin}in',
            },
            print_background=True,
            prefer_css_page_size=False,
        )
        await context.close()

        # Return as downloadable file
        return Response(
            pdf_bytes,
            mimetype='application/pdf',
            headers={
                'Content-Disposition': 'attachment; filename=website.pdf',
                'Content-Length': str(len(pdf_bytes))
            }
        )

    except PlaywrightTimeoutError:
        return jsonify({'error': 'Timeout: The page did not load in time.'}), 408
    except Exception as e:
        return jsonify({'error': f'PDF generation failed: {str(e)}'}), 500


def is_valid_url(url):
    """Basic URL validation."""
    try:
        result = urlparse(url)
        return all([result.scheme, result.netloc])
    except:
        return False


@app.teardown_appcontext
async def close_browser(exception=None):
    """Clean up browser resources on app shutdown."""
    global _browser, _playwright
    if _browser:
        await _browser.close()
        _browser = None
    if _playwright:
        await _playwright.stop()
        _playwright = None


if __name__ == '__main__':
    # For local development; on production use gunicorn or similar.
    app.run(debug=True, host='0.0.0.0', port=5000)